|
|
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_workflow_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}") |
|
|
|
|
|
|
|
|
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)) |
