diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..aff3050 --- /dev/null +++ b/agent.py @@ -0,0 +1,133 @@ +import os +from dotenv import load_dotenv +from langchain_ollama import OllamaLLM, OllamaEmbeddings +from langchain_chroma import Chroma +from langchain.text_splitter import RecursiveCharacterTextSplitter +from langchain_community.document_loaders import DirectoryLoader, TextLoader, UnstructuredMarkdownLoader +from langchain.tools import Tool +from langchain_community.tools.tavily_search import TavilySearchResults +from langchain.agents import initialize_agent, AgentType + +# Load environment variables +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") + +# Configuration +PERSIST_DIRECTORY = "./chroma_db" +DOCUMENTS_DIR = "./documents" +EMBEDDING_MODEL = "nomic-embed-text" +LLM_MODEL = "llama3" + +def setup_vectorstore(): + """Initialize or load ChromaDB vectorstore with Ollama embeddings.""" + embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) + + if os.path.exists(PERSIST_DIRECTORY) and os.listdir(PERSIST_DIRECTORY): + vectorstore = Chroma( + persist_directory=PERSIST_DIRECTORY, + embedding_function=embeddings + ) + else: + # Load and process documents + loader = DirectoryLoader( + DOCUMENTS_DIR, + glob="**/*", + loader_cls=lambda path: TextLoader(path, encoding="utf-8") if path.endswith(".txt") + else UnstructuredMarkdownLoader(path) if path.endswith(".md") + else None, + show_progress=True, + use_multithreading=True + ) + documents = loader.load() + + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, + chunk_overlap=200, + length_function=len + ) + texts = text_splitter.split_documents(documents) + + vectorstore = Chroma.from_documents( + documents=texts, + embedding=embeddings, + persist_directory=PERSIST_DIRECTORY + ) + vectorstore.persist() + + return vectorstore + +# Initialize vectorstore +vectorstore = setup_vectorstore() + +# Define tools +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]) + +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]) + +# Create LangChain tools +local_tool = Tool( + name="search_local_kb", + func=search_local_kb, + description="Useful for answering questions about local documents stored in the knowledge base." +) + +web_tool = Tool( + name="web_search", + func=web_search, + description="Useful for answering questions about current events, news, or general knowledge from the internet." +) + +tools = [local_tool, web_tool] + +# Initialize LLM +llm = OllamaLLM(model=LLM_MODEL) + +# Create agent with routing logic +agent = initialize_agent( + tools=tools, + llm=llm, + agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, + verbose=True, + handle_parsing_errors=True, + system_message="""You are an AI agent that must choose between two tools: + 1. search_local_kb: For questions about local documents (concepts, notes, stored information) + 2. web_search: For questions requiring up-to-date information from the internet + + When answering, ALWAYS specify your information source: + - If using local knowledge base: [Source: chromadb] + - If using web search: [Source: tavily] + + Be concise and accurate in your responses.""" +) + +# Chat interface +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") + except Exception as e: + print(f"Error: {str(e)}\n") + +if __name__ == "__main__": + main() \ No newline at end of file