'''
**simple agent demo app**

Before running, the user must set the following environment variables; otherwise, runtime exceptions will inevitably occur:
- MODEL_AGENT_NAME  # model id in Volcano Engine Ark platform
- MODEL_AGENT_API_KEY  # model api key in Volcano Engine Ark platform

MODEL_AGENT_NAME and MODEL_AGENT_API_KEY are used to access the model service of the Volcano Engine Ark platform.
'''
import logging
import os

from veadk import Agent, Runner
from veadk.knowledgebase import KnowledgeBase


from agentkit.apps import AgentkitSimpleApp

logger = logging.getLogger(__name__)

collection_name = os.getenv("DATABASE_VIKING_COLLECTION")
if not collection_name:
    raise ValueError("DATABASE_VIKING_COLLECTION environment variable is not set")

model_name = "doubao-seed-1-6-250615"

app = AgentkitSimpleApp()
knowledgebase = KnowledgeBase(backend="viking", index=collection_name)


agent = Agent(
    instruction="Answer customer's questions according to your knowledgebase.",
    model_name=model_name,
    knowledgebase=knowledgebase
)
runner = Runner(agent=agent)


@app.entrypoint
async def run(payload: dict, headers: dict) -> str:
    prompt = payload["prompt"]
    user_id = headers["user_id"]
    session_id = headers["session_id"]

    logger.info(
        f"Running agent with prompt: {prompt}, user_id: {user_id}, session_id: {session_id}"
    )
    response = await runner.run(messages=prompt, user_id=user_id, session_id=session_id)

    logger.info(f"Run response: {response}")
    return response


@app.ping
def ping() -> str:
    return "pong!"


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000)