feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

src/index.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+RAG Agent using ChromaDB for vector storage and Tavily for web search.
+
+The agent can ingest web pages (or arbitrary text) into a Chroma collection
+and answer queries by retrieving relevant documents and passing them to an
+OpenAI LLM.
+
+Prerequisites:
+- OpenAI API key set in the environment variable OPENAI_API_KEY
+- Tavily API key set in the environment variable TAVILY_API_KEY
+- Python 3.9+
+
+Usage:
+    python src/index.py ingest <url_or_text>   # Ingest a URL or raw text
+    python src/index.py query <question>        # Query the agent
+
+Example:
+    python src/index.py ingest https://en.wikipedia.org/wiki/OpenAI
+    python src/index.py query "What is OpenAI?"
+"""
+
+import os
+import sys
+import argparse
+from pathlib import Path
+from typing import List
+
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.llms import OpenAI
+from langchain.chains import RetrievalQA
+from tavily import TavilyClient
+
+# --------------------------------------------------------------------------- #
+# Configuration
+# --------------------------------------------------------------------------- #
+
+# Directory where the ChromaDB data will be stored
+CHROMA_DATA_DIR = Path.home() / ".rag_agent" / "chromadb"
+
+# Name of the collection used for storing documents
+COLLECTION_NAME = "rag_collection"
+
+# --------------------------------------------------------------------------- #
+# Helper functions
+# --------------------------------------------------------------------------- #
+
+def get_chroma_collection() -> Chroma:
+    """
+    Create or load a Chroma collection.
+    """
+    embeddings = OpenAIEmbeddings()
+    return Chroma(
+        collection_name=COLLECTION_NAME,
+        embedding_function=embeddings,
+        persist_directory=str(CHROMA_DATA_DIR),
+    )
+
+def ingest_text(text: str, collection: Chroma) -> None:
+    """
+    Add raw text to the Chroma collection.
+    """
+    collection.add_texts([text])
+
+def ingest_url(url: str, collection: Chroma, tavily_client: TavilyClient) -> None:
+    """
+    Fetch content from a URL using Tavily, embed it, and store it in Chroma.
+    """
+    # Tavily's search returns a list of results; we use the first result's content.
+    results = tavily_client.search(query=url, max_results=1)
+    if not results:
+        print(f"No results found for URL: {url}")
+        return
+    content = results[0].content
+    if not content:
+        print(f"No content extracted from URL: {url}")
+        return
+    collection.add_texts([content])
+    print(f"Ingested content from {url}")
+
+def query_agent(question: str, collection: Chroma) -> str:
+    """
+    Retrieve relevant documents from Chroma and ask OpenAI to answer.
+    """
+    llm = OpenAI(temperature=0)
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=collection.as_retriever(search_kwargs={"k": 4}),
+    )
+    return qa_chain.run(question)
+
+# --------------------------------------------------------------------------- #
+# Main entry point
+# --------------------------------------------------------------------------- #
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="RAG Agent with ChromaDB & Tavily")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    ingest_parser = subparsers.add_parser("ingest", help="Ingest a URL or raw text")
+    ingest_parser.add_argument("source", help="URL or raw text to ingest")
+
+    query_parser = subparsers.add_parser("query", help="Ask a question")
+    query_parser.add_argument("question", help="The question to ask the agent")
+
+    args = parser.parse_args()
+
+    # Ensure required environment variables are set
+    if "OPENAI_API_KEY" not in os.environ:
+        print("Error: OPENAI_API_KEY environment variable not set.")
+        sys.exit(1)
+    if "TAVILY_API_KEY" not in os.environ:
+        print("Error: TAVILY_API_KEY environment variable not set.")
+        sys.exit(1)
+
+    # Initialize Chroma collection
+    collection = get_chroma_collection()
+
+    # Initialize Tavily client
+    tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+
+    if args.command == "ingest":
+        source = args.source
+        if source.startswith(("http://", "https://")):
+            ingest_url(source, collection, tavily_client)
+        else:
+            ingest_text(source, collection)
+            print("Ingested raw text.")
+    elif args.command == "query":
+        answer = query_agent(args.question, collection)
+        print("\nAnswer:\n")
+        print(answer)
+
+if __name__ == "__main__":
+    main()