diff --git a/agent.py b/agent.py index 112085c..f411e57 100644 --- a/agent.py +++ b/agent.py @@ -63,15 +63,15 @@ agent = initialize_agent( def main(): print("RAG Agent with ChromaDB and Tavily Web Search") print("Type 'exit' to quit\n") - + while True: query = input("Запрос: ").strip() if query.lower() == "exit": break - + if not query: continue - + try: response = agent.run(query) print(f"\n{response}\n") @@ -79,4 +79,56 @@ def main(): print(f"Error: {str(e)}\n") if __name__ == "__main__": - main() \ No newline at end of file + main() +``` + +```python +# main.py + +import os +from dotenv import load_dotenv +from langchain_community.tools.tavily_search import TavilySearchResults +from langchain.tools import tool +from vectorstore import get_vectorstore + +load_dotenv() +TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") +if not TAVILY_API_KEY: + raise ValueError("TAVILY_API_KEY not found in .env file") + +vectorstore = get_vectorstore() + +@tool +def search_local_kb(query: str) -> str: + """Search local knowledge base using ChromaDB.""" + retriever = vectorstore.as_retriever(search_kwargs={"k": 3}) + docs = retriever.get_relevant_documents(query) + return "\n\n".join([doc.page_content for doc in docs]) + +@tool +def web_search(query: str) -> str: + """Search the web using Tavily.""" + search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3) + results = search.run(query) + return "\n\n".join([result["content"] for result in results]) +``` + +```python +# tools.py + +import os +from dotenv import load_dotenv +from langchain_community.tools.tavily_search import TavilySearchResults +from langchain.tools import tool + +load_dotenv() +TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") +if not TAVILY_API_KEY: + raise ValueError("TAVILY_API_KEY not found in .env file") + +@tool +def web_search(query: str) -> str: + """Search the web using Tavily.""" + search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3) + results = search.run(query) + return "\n\n".join([result["content"] for result in results]) \ No newline at end of file