| from src.state.state import State |
| from src.tools.websearch import WebSearchTool |
| from langchain_core.messages import HumanMessage, AIMessage, ToolMessage |
| from langchain.tools import BaseTool |
|
|
| class WebSearchChatbot: |
| def __init__(self, model, session_id: str = "default", tavily_api_key: str = None): |
| self.model = model |
| self.session_id = session_id |
| self.memory_config = {"configurable": {"session_id": session_id}} |
| |
| if tavily_api_key and tavily_api_key.strip(): |
| try: |
| self.web_search = WebSearchTool(tavily_api_key) |
| self.tools = [self.web_search.get_tool()] |
| self.model_with_tools = model.bind_tools(self.tools) |
| self.has_search = True |
| except Exception as e: |
| self.model_with_tools = model |
| self.has_search = False |
| else: |
| self.model_with_tools = model |
| self.has_search = False |
|
|
| def process(self, state): |
| messages = state['messages'] |
| if not messages: |
| return state |
| |
| if not self.has_search: |
| |
| last_message = messages[-1] |
| if hasattr(last_message, 'content') and any(keyword in last_message.content.lower() for keyword in ['search', 'find', 'latest', 'current', 'news']): |
| search_disclaimer = "I don't have web search capabilities enabled. Please provide a Tavily API key to search for current information." |
| response_content = f"{search_disclaimer}\n\nBased on my training data, I can still help with general questions." |
| from langchain_core.messages import AIMessage |
| return {'messages': AIMessage(content=response_content)} |
| |
| response = self.model_with_tools.invoke(messages, config=self.memory_config) |
| |
| if hasattr(response, 'tool_calls') and response.tool_calls: |
| messages.append(response) |
| |
| for tool_call in response.tool_calls: |
| tool_result = self._execute_tool_call(tool_call) |
| tool_message = ToolMessage( |
| content=str(tool_result), |
| tool_call_id=tool_call['id'] |
| ) |
| messages.append(tool_message) |
| |
| final_response = self.model_with_tools.invoke(messages, config=self.memory_config) |
| return {'messages': final_response} |
| |
| return {'messages': response} |
| |
| def _execute_tool_call(self, tool_call): |
| if tool_call['name'] == 'tavily_search_results_json': |
| return self.web_search.search_tool.invoke(tool_call['args']) |
| return "Tool not found" |
|
|
