diff --git a/main.py b/main.py index 75fb450..96f4768 100644 --- a/main.py +++ b/main.py @@ -1,58 +1,152 @@ import os +import sys import asyncio -from dotenv import load_dotenv -from langchain_openai import ChatOpenAI +from pathlib import Path +from typing import List + +from langchain_ollama import Ollama, OllamaEmbeddings +from langchain_core.documents import Document +from langchain_qdrant import QdrantVectorStore +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.tools import tool from langchain_core.messages import HumanMessage from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from tools import search_knowledge_base, add_to_knowledge_base +from qdrant_client import QdrantClient -load_dotenv() +# Инициализация эмбеддингов и LLM через Ollama +embeddings = OllamaEmbeddings(model="nomic-embed-text") +llm = Ollama(model="llama3") -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=0.0, +# Инициализация Qdrant +client = QdrantClient(host="localhost", port=6333) +collection_name = "knowledge" +vector_store = QdrantVectorStore( + client=client, + collection_name=collection_name, + embedding=embeddings, ) -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) +# Чанкинг +splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + +# Инструмент: поиск в базе знаний +@tool +def search_knowledge_base(query: str, max_results: int = 3) -> str: + """Search the knowledge base for relevant information.""" + docs: List[Document] = vector_store.similarity_search(query, k=max_results) + if not docs: + return "No results." + return "\n".join(doc.page_content for doc in docs) + +# Инструмент: добавление документа в базу знаний +@tool +def add_to_knowledge_base(content: str, title: str = "doc") -> str: + """Add content to the knowledge base.""" + chunks = splitter.split_text(content) + docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] + vector_store.add_documents(docs) + return f"Added: {title} ({len(chunks)} chunks)." + +# Backend для deepagents +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 information. " + "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=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt="You are a helpful knowledge assistant. Use the tools to search and add documents.", + system_prompt=system_prompt, ) -async def interactive_loop(): - thread_id = "interactive-session" - print("Welcome to RAG Agent. Commands: /add