from llama_index.core.agent.workflow import FunctionAgent, AgentWorkflow
from llama_index.llms.vllm import Vllm

from prompts import (
    ORCHESTRATOR_SYSTEM_PROMPT,
    NOTION_SYSTEM_PROMPT,
    IMPLEMENTATION_SYSTEM_PROMPT,
)
from tools import notion_retrieval_tool, implementation_tool

llm = Vllm(
    model="model_name",
    tensor_parallel_size=4,
    max_new_tokens=100,
    vllm_kwargs={"swap_space": 1, "gpu_memory_utilization": 0.5},
)


orchestrator_agent = FunctionAgent(
    name="OrchestratorAgent",
    description=(
        "You are the OrchestratorAgent responsible for coordinating tasks between multiple agents. "
     
    ),
    system_prompt=ORCHESTRATOR_SYSTEM_PROMPT,
    llm=llm,
    tools=[],
    can_handoff_to=["NotionAgent", "ImplementationAgent"],
)

  parser = LlamaParse(
            api_key=LLAMAINDEX_API_KEY,
            parsing_instruction=parsing_instruction,
            premium_mode=True,
            split_by_page=False,
            verbose=False
        )
        file_extractor = {".pdf": parser, ".docx": parser, ".doc": parser}
        documents = SimpleDirectoryReader(
            input_files=[file_path], file_extractor=file_extractor
        ).load_data()

        document_text = "\n".join(
            [doc.text for doc in documents if hasattr(doc, 'text')]
        )