diff --git a/main.py b/main.py index f67d0c0..75fb450 100644 --- a/main.py +++ b/main.py @@ -1,20 +1,14 @@ import os import asyncio from dotenv import load_dotenv -from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage -from langchain_core.documents import Document -from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_qdrant import QdrantVectorStore -from langchain_text_splitters import RecursiveCharacterTextSplitter -from qdrant_client import QdrantClient +from tools import search_knowledge_base, add_to_knowledge_base -# Загрузка переменных окружения load_dotenv() -# Инициализация LLM llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -22,90 +16,43 @@ llm = ChatOpenAI( temperature=0.0, ) -# Инициализация эмбеддингов -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) - -# Инициализация Qdrant -client = QdrantClient(url="http://localhost:6333") -collection_name = "knowledge_base" -vector_store = QdrantVectorStore( - client=client, - collection_name=collection_name, - embeddings=embeddings, -) - -# Чанкинг -splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - -# Инструмент поиска -@tool -def search_knowledge_base(query: str, max_results: int) -> str: - """Semantic search in the knowledge base.""" - docs = vector_store.similarity_search(query, k=max_results) - if not docs: - return "No results found." - return "\n".join(doc.page_content for doc in docs) - -# Инструмент добавления -@tool -def add_to_knowledge_base(content: str, title: str) -> 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 {len(docs)} chunks for {title}." - -# Backend для deepagents backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# Создание агента agent = create_deep_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt="You are a helpful agent with access to a knowledge base. Use the tools to search and add information.", + system_prompt="You are a helpful knowledge assistant. Use the tools to search and add documents.", ) -async def main(): - print("Interactive agent. Commands: /add, /search, /quit") +async def interactive_loop(): + thread_id = "interactive-session" + print("Welcome to RAG Agent. Commands: /add