Удалить agent.py

2026-05-28 13:32:30 +00:00
parent a42faf40d1
commit a26dcf2dc2
1 changed files with 0 additions and 140 deletions
@@ -1,140 +0,0 @@
 """
 Agent and vector store setup for the RAG task.
 This module defines:
 * `QdrantVectorStore` wrapper that uses Ollama embeddings.
 * Two tools – ``search_knowledge_base`` and ``add_to_knowledge_base``.
 * A helper to create the agent via :func:`langchain.agents.create_agent`.
 """
 import os
 from pathlib import Path
 from typing import List, Dict
 from langchain_ollama import ChatOllama, OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain.tools import tool
 from langchain.agents import create_agent
 from langchain_core.messages import HumanMessage
 # ---------------------------------------------------------------------------
 # Vector store configuration
 # ---------------------------------------------------------------------------
 QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
 QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
 COLLECTION_NAME = "knowledge"
 # Initialize embeddings and vector store.  The client is created lazily on first use.
 embeddings = OllamaEmbeddings(model="nomic-embed-text")
 vector_store: QdrantVectorStore | None = None
 def get_vector_store() -> QdrantVectorStore:
    """Return a singleton Qdrant vector store instance.
    The collection is created automatically if it does not exist.
    """
    global vector_store
    if vector_store is None:
        from qdrant_client import QdrantClient
        client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
        vector_store = QdrantVectorStore(
            client=client,
            collection_name=COLLECTION_NAME,
            embedding=embeddings,
        )
    return vector_store
 # ---------------------------------------------------------------------------
 # Tools
 # ---------------------------------------------------------------------------
@tool("search_knowledge_base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Semantic search in the knowledge base.
    Parameters
    ----------
    query: str
        Search query.
    max_results: int, optional
        Number of top results to return. Defaults to 5.
    Returns
    -------
    str
        A numbered list of passages or a message if nothing was found.
    """
    store = get_vector_store()
    docs = store.similarity_search(query, k=max_results)
    if not docs:
        return "No relevant documents found."
    return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
@tool("add_to_knowledge_base")
 def add_to_knowledge_base(content: str, title: str = "document") -> str:
    """Add a document to the knowledge base.
    Parameters
    ----------
    content: str
        Text of the document.
    title: str, optional
        Title used as metadata. Defaults to ``"document"``.
    Returns
    -------
    str
        Confirmation message.
    """
    store = get_vector_store()
    from langchain_core.documents import Document
    doc = Document(page_content=content, metadata={"title": title})
    store.add_documents([doc])
    return f"Added '{title}' to knowledge base."
 # ---------------------------------------------------------------------------
 # Agent creation helper
 # ---------------------------------------------------------------------------
 llm = ChatOllama(model="llama3", temperature=0.0)
 SYSTEM_PROMPT = (
    "You are an assistant that can search and add information to a local knowledge base.\n"
    "Use the tools `search_knowledge_base` and `add_to_knowledge_base`.\n"
    "When searching, return the most relevant passages. When adding, confirm success."
 )
 def create_rag_agent():
    """Return a LangChain agent configured with the RAG tools."""
    agent = create_agent(
        llm=llm,
        tools=[search_knowledge_base, add_to_knowledge_base],
        system_prompt=SYSTEM_PROMPT,
    )
    return agent
 # ---------------------------------------------------------------------------
 # Example usage (can be imported by main.py)
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
    # Simple demo: add a short doc and search it.
    agent = create_rag_agent()
    print("Adding sample document...")
    res = agent.invoke(
        {"messages": [HumanMessage(content="Add to knowledge base: content='Python is great' title='Python intro'")]},
        {"configurable": {"thread_id": "demo-1"}},
    )
    print(res["messages"][-1].content)
    print("Searching for Python...")
    res = agent.invoke(
        {"messages": [HumanMessage(content="Search for Python')"],},
        {"configurable": {"thread_id": "demo-2"}},
    )
    print(res["messages"][-1].content)
 # End of agent.py