cli_research.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. import argparse
  2. from langchain_core.messages import HumanMessage
  3. from agent.graph import graph
  4. def main() -> None:
  5. """Run the research agent from the command line."""
  6. parser = argparse.ArgumentParser(description="Run the LangGraph research agent")
  7. parser.add_argument("question", help="Research question")
  8. parser.add_argument(
  9. "--initial-queries",
  10. type=int,
  11. default=3,
  12. help="Number of initial search queries",
  13. )
  14. parser.add_argument(
  15. "--max-loops",
  16. type=int,
  17. default=2,
  18. help="Maximum number of research loops",
  19. )
  20. parser.add_argument(
  21. "--reasoning-model",
  22. default="gemini-2.5-pro-preview-05-06",
  23. help="Model for the final answer",
  24. )
  25. args = parser.parse_args()
  26. state = {
  27. "messages": [HumanMessage(content=args.question)],
  28. "initial_search_query_count": args.initial_queries,
  29. "max_research_loops": args.max_loops,
  30. "reasoning_model": args.reasoning_model,
  31. }
  32. result = graph.invoke(state)
  33. messages = result.get("messages", [])
  34. if messages:
  35. print(messages[-1].content)
  36. if __name__ == "__main__":
  37. main()