Удалить agent.py
This commit is contained in:
@@ -1,79 +0,0 @@
|
|||||||
"""
|
|
||||||
Agent definition for the RAG system.
|
|
||||||
|
|
||||||
Provides two tools:
|
|
||||||
* search_knowledge_base(query, max_results)
|
|
||||||
* add_to_knowledge_base(content, title)
|
|
||||||
|
|
||||||
The agent is created with create_agent from langchain.agents.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
|
||||||
from typing import List
|
|
||||||
from langchain_ollama import OllamaEmbeddings
|
|
||||||
from chromadb import PersistentClient
|
|
||||||
from chromadb.utils import embedding_functions as ef
|
|
||||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
||||||
from langchain.tools import tool
|
|
||||||
from langchain.agents import create_agent
|
|
||||||
from langchain_core.messages import HumanMessage
|
|
||||||
|
|
||||||
# Initialize embeddings and Chroma client
|
|
||||||
EMBEDDINGS = OllamaEmbeddings(model="nomic-embed-text")
|
|
||||||
CHROMA_PATH = os.path.join(os.getcwd(), "chromadb_store")
|
|
||||||
CLIENT = PersistentClient(path=CHROMA_PATH)
|
|
||||||
COLLECTION_NAME = "knowledge"
|
|
||||||
if COLLECTION_NAME not in CLIENT.list_collections():
|
|
||||||
CLIENT.create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
|
|
||||||
COLL = CLIENT.get_or_create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
|
|
||||||
|
|
||||||
# Chunker from chunker.py
|
|
||||||
from chunker import CHUNKER
|
|
||||||
|
|
||||||
@tool
|
|
||||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
||||||
"""
|
|
||||||
Add a document to the knowledge base.
|
|
||||||
The content is split into chunks and stored with metadata.
|
|
||||||
Returns confirmation message.
|
|
||||||
"""
|
|
||||||
# Split content
|
|
||||||
chunks = CHUNKER.split_text(content)
|
|
||||||
ids = [f"{title}_{i}" for i in range(len(chunks))]
|
|
||||||
metadatas = [{"title": title} for _ in chunks]
|
|
||||||
COLL.add(ids=ids, documents=chunks, metadatas=metadatas)
|
|
||||||
return f"Added {len(chunks)} chunks from '{title}'."
|
|
||||||
|
|
||||||
@tool
|
|
||||||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
||||||
"""
|
|
||||||
Search the knowledge base for relevant documents.
|
|
||||||
Returns a formatted string of results.
|
|
||||||
"""
|
|
||||||
results = COLL.query(
|
|
||||||
query_texts=[query],
|
|
||||||
n_results=max_results,
|
|
||||||
include=['documents', 'distances'],
|
|
||||||
)
|
|
||||||
docs = results.get("documents", [])[0]
|
|
||||||
dists = results.get("distances", [])[0]
|
|
||||||
if not docs:
|
|
||||||
return "No relevant documents found."
|
|
||||||
output_lines = []
|
|
||||||
for i, (doc, dist) in enumerate(zip(docs, dists), 1):
|
|
||||||
output_lines.append(f"{i}. (score: {dist:.4f})\n{doc[:200]}...")
|
|
||||||
return "\n\n".join(output_lines)
|
|
||||||
|
|
||||||
# Create agent
|
|
||||||
SYSTEM_PROMPT = (
|
|
||||||
"You are an assistant that can search and add to a knowledge base. Use the provided tools."
|
|
||||||
)
|
|
||||||
AGENT = create_agent(
|
|
||||||
llm=None, # No LLM needed for tool calls; agent will use system prompt only
|
|
||||||
tools=[add_to_knowledge_base, search_knowledge_base],
|
|
||||||
system_prompt=SYSTEM_PROMPT,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Expose a simple invoke function
|
|
||||||
async def run_agent(messages: List[HumanMessage]):
|
|
||||||
return await AGENT.ainvoke({"messages": messages}, {"configurable": {"thread_id": "rag-agent"}})
|
|
||||||
Reference in New Issue
Block a user