agent.py

from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from tools.simple_tools import (
    search_web_tool, record_notes_tool, write_report_tool, review_report_tool,
    get_workflow_state, reset_workflow_state
)
from dotenv import load_dotenv
import os
from llama_index.core.agent.workflow import AgentWorkflow, ReActAgent
from llama_index.core.workflow import Context

load_dotenv(os.path.join(os.path.dirname(__file__), 'env.local'))

class LlamaIndexReportAgent:
    def __init__(self):
        self.llm = HuggingFaceInferenceAPI(
            model_name="microsoft/Phi-3.5-mini-instruct",
            token=os.getenv("HUGGING_FACE_TOKEN")
        )

        self.research_agent = ReActAgent(
            name="ResearchAgent",
            description="Searches the web and records notes.",
            system_prompt=(
                "You are a Research Agent. Your ONLY job is to research and hand off to WriteAgent.\n"
                "\n"
                "STRICT WORKFLOW:\n"
                "1. Use search_web tool to search for information\n"
                "2. Use record_notes tool to save what you found\n"
                "3. Say: 'Research complete. I have gathered sufficient information. Handing off to WriteAgent.'\n"
                "\n"
                "CRITICAL RULES:\n"
                "- You can ONLY use search_web and record_notes tools\n"
                "- You CANNOT write reports - that's WriteAgent's job\n"
                "- You CANNOT use write_report tool - you don't have access to it\n"
                "- After research, you MUST hand off with the exact message above\n"
                "- Do NOT attempt to write any report content yourself\n"
                "\n"
                "AVAILABLE TOOLS: search_web, record_notes\n"
                "HANDOFF MESSAGE: 'Research complete. I have gathered sufficient information. Handing off to WriteAgent.'"
            ),
            tools=[search_web_tool, record_notes_tool],
            llm=self.llm,
            can_handoff_to=["WriteAgent"],
        )

        self.write_agent = ReActAgent(
            name="WriteAgent",
            description="Writes a structured report based on research notes.",
            system_prompt=(
                "You are a Writing Agent. Your purpose is to create a concise, well-structured report.\n"
                "\n"
                "INSTRUCTIONS:\n"
                "1. Check if there's any feedback from ReviewAgent (not 'Review required.')\n"
                "2. If there's feedback, revise the report accordingly\n"
                "3. If no feedback, create initial report based on research\n"
                "4. MUST call write_report tool with these parameters:\n"
                "   - report_content: Concise markdown report (200-400 words)\n"
                "   - title: Descriptive report title\n"
                "5. Report structure (keep sections brief):\n"
                "   - # Main Title\n"
                "   - ## Introduction (1-2 sentences)\n"
                "   - ## Key Points (2-3 bullet points)\n"
                "   - ## Conclusion (1-2 sentences)\n"
                "6. After calling tool: 'Report written. Handing off to ReviewAgent.'\n"
                "\n"
                "CRITICAL: Keep the report_content CONCISE to avoid truncation!\n"
                "You MUST actually call the write_report tool with proper parameters!"
            ),
            tools=[write_report_tool],
            llm=self.llm,
            can_handoff_to=["ReviewAgent"],
        )

        self.review_agent = ReActAgent(
            name="ReviewAgent",
            description="Reviews the written report.",
            system_prompt=(
                "You are a Reviewing Agent. Your purpose is to review the report quality.\n"
                "1. Check the report content that was written\n"
                "2. Use review_report tool to provide feedback\n"
                "3. If report is good quality, start feedback with 'APPROVED:'\n"
                "4. If needs improvement, provide specific suggestions and hand off to WriteAgent\n"
                "5. Quality criteria: clear structure, sufficient detail, proper formatting"
            ),
            tools=[review_report_tool],
            llm=self.llm,
            can_handoff_to=["WriteAgent"],
        )

        self.agent_workflow = AgentWorkflow(
            agents=[self.research_agent, self.write_agent, self.review_agent],
            root_agent=self.research_agent.name,
            initial_state={
                "research_notes": {},
                "report_content": "Not written yet.",
                "review": "Review required.",
            },
        )

    def get_final_state(self) -> dict:
        """Get the final workflow state from the simple tools."""
        return get_workflow_state()

    async def run_workflow(self, user_msg=None):
        if user_msg is None:
            user_msg = (
                "Write me a report on the history of the internet. "
                "Briefly describe the history of the internet, including the development of the internet, the development of the web, "
                "and the development of the internet in the 21st century."
            )
        
        # Reset state for new workflow
        reset_workflow_state()
        
        # Create context and initialize state
        ctx = Context(self.agent_workflow)
        await ctx.set("state", {
            "research_notes": {},
            "report_content": "Not written yet.",
            "review": "Review required.",
        })
        
        handler = self.agent_workflow.run(user_msg=user_msg, ctx=ctx)

        current_agent = None
        async for event in handler.stream_events():
            if (
                hasattr(event, "current_agent_name")
                and event.current_agent_name != current_agent
            ):
                current_agent = event.current_agent_name
                print(f"\n{'='*50}")
                print(f"🤖 Agent: {current_agent}")
                print(f"{'='*50}\n")

            if hasattr(event, "response") and hasattr(event.response, "content"):
                if event.response.content:
                    print("📤 Output:", event.response.content)
                if hasattr(event, "tool_calls") and event.tool_calls:
                    print(
                        "🛠️  Planning to use tools:",
                        [call.tool_name for call in event.tool_calls],
                    )
            elif hasattr(event, "tool_name") and hasattr(event, "tool_output"):
                print(f"🔧 Tool Result ({event.tool_name}):")
                print(f"  Arguments: {getattr(event, 'tool_kwargs', {})}")
                print(f"  Output: {event.tool_output}")
            elif hasattr(event, "tool_name") and hasattr(event, "tool_kwargs"):
                print(f"🔨 Calling Tool: {event.tool_name}")
                print(f"  With arguments: {event.tool_kwargs}")

        # After the workflow completes, print the final report
        final_state = self.get_final_state()
        print(f"\n📊 Final State:")
        print(f"Research notes: {len(final_state.get('research_notes', {}))}")
        print(f"Report written: {final_state.get('report_content', 'Not written') != 'Not written yet.'}")
        print(f"Review: {final_state.get('review', 'No review')[:100]}...")
        
        if final_state.get("structured_report"):
            print("\n📄 Final Report Generated Successfully!")
            report = final_state["structured_report"]
            print(f"Title: {report['title']}")
            print(f"Word count: {report['word_count']}")
            print(f"Sections: {len(report['sections'])}")
        else:
            print("\n⚠️ No final report was generated by the workflow.")

if __name__ == "__main__":
    import asyncio
    agent = LlamaIndexReportAgent()
    user_msg = input("Enter the topic or instructions for the report (leave blank for default): ").strip()
    if not user_msg:
        user_msg = None
    asyncio.run(agent.run_workflow(user_msg=user_msg))