feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

agent.py

@@ -0,0 +1,99 @@
+import os
+from typing import Any, Dict
+
+from langchain_ollama import Ollama
+from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType
+from langchain.schema import AgentAction, AgentFinish
+from langchain.tools import tool
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers import StrOutputParser
+
+from vectorstore import QdrantVectorStore
+from tools import search_local_kb, web_search
+
+
+def create_agent(
+    vectorstore: QdrantVectorStore,
+    tavily_api_key: str,
+    llm_model: str = "llama3",
+) -> AgentExecutor:
+    """
+    Create a LangChain agent that routes queries to either the local KB or the web.
+
+    Parameters
+    ----------
+    vectorstore : QdrantVectorStore
+        The vector store for local knowledge base.
+    tavily_api_key : str
+        Tavily API key.
+    llm_model : str
+        Ollama model name for LLM.
+
+    Returns
+    -------
+    AgentExecutor
+        Configured agent executor.
+    """
+    # LLM
+    llm = Ollama(model=llm_model)
+
+    # Define tools with partial application of required arguments
+    tools = [
+        Tool(
+            name="search_local_kb",
+            func=lambda query, top_k=3: search_local_kb(
+                query=query, top_k=top_k, vectorstore=vectorstore
+            ),
+            description=(
+                "Use this tool to search the local knowledge base. "
+                "Return the most relevant snippets."
+            ),
+        ),
+        Tool(
+            name="web_search",
+            func=lambda query, max_results=3: web_search(
+                query=query, tavily_api_key=tavily_api_key, max_results=max_results
+            ),
+            description=(
+                "Use this tool to search the web via Tavily. "
+                "Return the most relevant snippets."
+            ),
+        ),
+    ]
+
+    # Prompt template for the agent
+    system_prompt = (
+        "You are an AI assistant that can answer questions using either "
+        "the local knowledge base or the web. If the question is about "
+        "recent events, news, or requires up-to-date information, "
+        "use the web_search tool. If the question is about "
+        "information contained in the local documents, use the "
+        "search_local_kb tool. After retrieving the information, "
+        "provide a concise answer and state the source (chromadb or tavily)."
+    )
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder("history"),
+            ("human", "{input}"),
+            MessagesPlaceholder("agent_scratchpad"),
+        ]
+    )
+
+    # Output parser
+    output_parser = StrOutputParser()
+
+    # Agent
+    agent = initialize_agent(
+        tools=tools,
+        llm=llm,
+        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=False,
+        handle_parsing_errors=True,
+        max_iterations=5,
+        early_stopping_method="generate",
+        system_message=system_prompt,
+    )
+
+    return agent