diff --git a/rag_agent.py b/rag_agent.py new file mode 100644 index 0000000..e78585c --- /dev/null +++ b/rag_agent.py @@ -0,0 +1,75 @@ +import os +from langchain_ollama import ChatOllama, OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_core.documents import Document +from langchain.tools import tool +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from deepagents import create_deep_agent + +# Настройки окружения +OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://localhost:11434") +QDRANT_HOST = os.getenv("QDRANT_HOST", "http://localhost:6333") + +# Эмбеддинги +embeddings = OllamaEmbeddings( + model="nomic-embed-text", + base_url=OLLAMA_HOST +) + +# Векторная база +vector_store = QdrantVectorStore( + url=QDRANT_HOST, + collection_name="knowledge", + embedding_function=embeddings +) + +# Чанкинг +chunker = RecursiveCharacterTextSplitter( + chunk_size=500, + chunk_overlap=50 +) + +@tool +def add_to_knowledge_base(content: str, title: str = "doc") -> str: + """Add content to the knowledge base.""" + chunks = chunker.split_text(content) + docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] + vector_store.add_documents(docs) + return f"Added: {title} with {len(chunks)} chunks." + +@tool +def search_knowledge_base(query: str, max_results: int = 3) -> str: + """Search the knowledge base for relevant information.""" + docs = vector_store.similarity_search(query, k=max_results) + if not docs: + return "No results." + return "\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs)) + +# LLM +llm = ChatOllama( + model="llama3", + base_url=OLLAMA_HOST +) + +# Backend +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +# Системный промпт +system_prompt = ( + "You are a helpful agent with access to a knowledge base. " + "Use the tools to search and add knowledge. " + "When you need to retrieve information, call search_knowledge_base. " + "When you need to store new information, call add_to_knowledge_base." +) + +# Агент +agent = create_deep_agent( + model=llm, + tools=[add_to_knowledge_base, search_knowledge_base], + backend=backend, + system_prompt=system_prompt, +) \ No newline at end of file