| 26 | error: Optional[str] |
| 27 | |
| 28 | class ScrapingAgent: |
| 29 | def __init__(self): |
| 30 | self.llm = get_llm() |
| 31 | self.llm_with_tools = self.llm.bind_tools(SCRAPING_TOOLS) |
| 32 | self.scraping_service = ScrapingService(settings.SCRAPEGRAPH_API_KEY) |
| 33 | self.graph = self._build_graph() |
| 34 | |
| 35 | def _build_graph(self) -> StateGraph: |
| 36 | graph = StateGraph(AgentState) |
| 37 | |
| 38 | # Add nodes |
| 39 | graph.add_node("analyze_request", self._analyze_request) |
| 40 | graph.add_node("call_tools", self._call_tools) |
| 41 | graph.add_node("generate_response", self._generate_response) |
| 42 | graph.add_node("execute_scraping", self._execute_scraping) |
| 43 | |
| 44 | # Set entry point |
| 45 | graph.set_entry_point("analyze_request") |
| 46 | |
| 47 | # Add edges |
| 48 | graph.add_conditional_edges( |
| 49 | "analyze_request", |
| 50 | self._should_use_tools, |
| 51 | { |
| 52 | True: "call_tools", |
| 53 | False: "generate_response" |
| 54 | } |
| 55 | ) |
| 56 | |
| 57 | graph.add_conditional_edges( |
| 58 | "call_tools", |
| 59 | self._should_execute_scraping, |
| 60 | { |
| 61 | True: "execute_scraping", |
| 62 | False: "generate_response" |
| 63 | } |
| 64 | ) |
| 65 | |
| 66 | graph.add_edge("execute_scraping", "generate_response") |
| 67 | graph.add_edge("generate_response", END) |
| 68 | |
| 69 | return graph.compile() |
| 70 | |
| 71 | async def _analyze_request(self, state: AgentState) -> AgentState: |
| 72 | """Analyze the user request and determine action.""" |
| 73 | messages = state["messages"] |
| 74 | |
| 75 | # Create prompt with system message |
| 76 | prompt = ChatPromptTemplate.from_messages([ |
| 77 | ("system", SYSTEM_PROMPT), |
| 78 | MessagesPlaceholder("messages") |
| 79 | ]) |
| 80 | |
| 81 | # Get response with tools |
| 82 | response = await self.llm_with_tools.ainvoke(messages) |
| 83 | |
| 84 | return { |
| 85 | **state, |