diff --git a/main.py b/main.py index 145da66..325f212 100644 --- a/main.py +++ b/main.py @@ -1,136 +1,19 @@ -import os import asyncio -from pathlib import Path - +import os from dotenv import load_dotenv -from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_chroma import Chroma -from langchain_core.documents import Document -from langchain.text_splitter import RecursiveCharacterTextSplitter -from langchain.tools import tool -from langchain_community.tools.tavily_search import TavilySearchResults - -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_core.messages import HumanMessage +from agent import run_agent load_dotenv() -# ---------- LLM ---------- -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, -) - -# ---------- Vector Store ---------- -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) - -def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - return Chroma( - collection_name="knowledge", - embedding_function=embeddings, - persist_directory=persist_directory, - ) - -vectorstore = create_vectorstore() - -def load_documents(directory: str, vectorstore: Chroma) -> None: - path = Path(directory) - if not path.is_dir(): - return - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs = [] - for file_path in path.rglob("*"): - if file_path.suffix.lower() not in {".txt", ".md"}: - continue - text = file_path.read_text(encoding="utf-8") - chunks = splitter.split_text(text) - for i, chunk in enumerate(chunks): - metadata = {"source": str(file_path), "chunk_index": i} - docs.append(Document(page_content=chunk, metadata=metadata)) - if docs: - vectorstore.add_documents(docs) - -# Load initial documents (run once or each start) -load_documents("./documents", vectorstore) - -# ---------- Tools ---------- -@tool -def search_local_kb(query: str, top_k: int = 3) -> str: - """ - Perform a semantic search in the local Chroma knowledge base. - Returns the concatenated contents of the most relevant documents. - """ - docs = vectorstore.similarity_search(query, k=top_k) - if not docs: - return "No relevant information found in the local knowledge base." - return "\n---\n".join(doc.page_content for doc in docs) - -tavily_tool = TavilySearchResults(max_results=5) - -@tool -def web_search(query: str) -> str: - """ - Search the web using Tavily and return a short summary of the top results. - """ - results = tavily_tool.run(query) - if not results: - return "No web results found." - # results is a list of dicts; extract title and url - lines = [] - for r in results: - title = r.get("title", "No title") - url = r.get("url", "") - lines.append(f"{title}: {url}") - return "\n".join(lines) - -# ---------- Backend ---------- -backend = CompositeBackend( - [ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), - ] -) - -# ---------- Agent ---------- -system_prompt = ( - "You are an AI assistant that can answer questions using two sources: " - "a local knowledge base (accessed via the tool `search_local_kb`) and the web (accessed via the tool `web_search`). " - "Decide which tool to use based on the user query. " - "When you use a tool, include the source name in your final answer: " - "`Source: chromadb` for local knowledge, `Source: tavily` for web results. " - "If both sources are needed, combine them and list both sources." -) - -agent = create_deep_agent( - model=llm, - tools=[search_local_kb, web_search], - backend=backend, - system_prompt=system_prompt, -) - -# ---------- CLI ---------- -async def chat_loop() -> None: - print("AI Assistant (type 'exit' to quit)") - thread_id = "session-1" +async def main(): + print("Введите запрос (или exit для выхода):") while True: - user_input = input("\nYou: ").strip() - if user_input.lower() in {"exit", "quit"}: - print("Goodbye!") + query = input("Запрос: ") + if query.lower() in ("exit", "quit"): + print("Выход.") break - - result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": thread_id}}, - ) - answer = result["messages"][-1].content - print(f"\nAssistant: {answer}") + response = await run_agent(query) + print(response) if __name__ == "__main__": - asyncio.run(chat_loop()) \ No newline at end of file + asyncio.run(main()) \ No newline at end of file