feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'

src/ingest.py

@@ -0,0 +1,72 @@
+"""
+Ingestion logic for FAQ documents into ChromaDB.
+"""
+
+import os
+from pathlib import Path
+from typing import List
+
+from chromadb import Client
+from chromadb.api.types import Documents, EmbeddingFunction
+from chromadb.config import Settings
+
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+def _load_faq_pairs(file_path: Path) -> List[tuple]:
+    """
+    Load FAQ pairs from a text file.
+    Expected format:
+        Q: <question>
+        A: <answer>
+    Each pair separated by a blank line.
+    """
+    pairs = []
+    with file_path.open("r", encoding="utf-8") as f:
+        content = f.read()
+    raw_pairs = content.strip().split("\n\n")
+    for raw in raw_pairs:
+        lines = raw.strip().splitlines()
+        if len(lines) < 2:
+            continue
+        q_line = lines[0].strip()
+        a_line = lines[1].strip()
+        if q_line.lower().startswith("q:") and a_line.lower().startswith("a:"):
+            question = q_line[2:].strip()
+            answer = a_line[2:].strip()
+            pairs.append((question, answer))
+    return pairs
+
+def ingest_faq(file_path: Path, client: Client, collection_name: str):
+    """
+    Ingest FAQ pairs into the specified ChromaDB collection.
+    """
+    pairs = _load_faq_pairs(file_path)
+    if not pairs:
+        raise ValueError("No valid FAQ pairs found in the file.")
+
+    # Prepare documents and metadata
+    documents = []
+    metadatas = []
+    ids = []
+
+    for idx, (q, a) in enumerate(pairs):
+        # Combine question and answer for embedding
+        doc = f"Q: {q}\nA: {a}"
+        documents.append(doc)
+        metadatas.append({"question": q, "answer": a})
+        ids.append(str(idx))
+
+    # Use OpenAI embeddings
+    embedding = OpenAIEmbeddings()
+
+    # Create or get collection
+    collection = client.get_or_create_collection(name=collection_name)
+
+    # Add documents to collection
+    collection.add(
+        documents=documents,
+        metadatas=metadatas,
+        ids=ids,
+        embedding_function=embedding
+    )

src/main.py

@@ -1,41 +1,67 @@
+#!/usr/bin/env python3
 """
-Command‑line interface for the FAQ bot.
+FAQ Bot using ChromaDB and LangChain
 """
 
-import argparse
 import os
+import sys
+import argparse
+from pathlib import Path
 
-from .bot import FAQBot
+from chromadb import Client
+from chromadb.config import Settings
+
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.llms.openai import OpenAIChat
+from langchain.chains import RetrievalQA
+from langchain.vectorstores import Chroma
+
+from ingest import ingest_faq
+from retriever import get_answer
+
+def init_chroma(collection_name: str) -> Client:
+    """
+    Initialize a ChromaDB client and create a collection if it does not exist.
+    """
+    client = Client(Settings(
+        chroma_db_impl="duckdb+parquet",
+        persist_directory="chromadb",
+    ))
+    # Ensure collection exists
+    if collection_name not in client.list_collections():
+        client.create_collection(name=collection_name)
+    return client
 
 def main():
     parser = argparse.ArgumentParser(description="FAQ Bot CLI")
-    parser.add_argument(
-        "--persist-dir",
-        type=str,
-        default="chromadb_persist",
-        help="Directory to persist ChromaDB data",
-    )
-    parser.add_argument(
-        "--openai-key",
-        type=str,
-        default=os.getenv("OPENAI_API_KEY"),
-        help="OpenAI API key (optional)",
-    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    ingest_parser = subparsers.add_parser("ingest", help="Ingest FAQ file into ChromaDB")
+    ingest_parser.add_argument("faq_file", type=Path, help="Path to FAQ text file")
+    ingest_parser.add_argument("--collection", type=str, default="faq_collection", help="Chroma collection name")
+
+    query_parser = subparsers.add_parser("ask", help="Ask a question to the FAQ bot")
+    query_parser.add_argument("question", type=str, help="Your question")
+    query_parser.add_argument("--collection", type=str, default="faq_collection", help="Chroma collection name")
+
     args = parser.parse_args()
 
-    bot = FAQBot(persist_dir=args.persist_dir, openai_api_key=args.openai_key)
+    # Ensure OpenAI API key is set
+    if "OPENAI_API_KEY" not in os.environ:
+        print("Error: OPENAI_API_KEY environment variable not set.", file=sys.stderr)
+        sys.exit(1)
 
-    print("FAQ Bot is ready. Type your question (Ctrl+C to exit).")
-    while True:
-        try:
-            question = input("\n> ")
-            if not question.strip():
-                continue
-            answer = bot.ask(question)
-            print(f"\nAnswer: {answer}")
-        except (KeyboardInterrupt, EOFError):
-            print("\nGoodbye!")
-            break
+    client = init_chroma(args.collection)
+
+    if args.command == "ingest":
+        ingest_faq(args.faq_file, client, args.collection)
+        print(f"Ingestion completed. Collection '{args.collection}' updated.")
+    elif args.command == "ask":
+        answer = get_answer(args.question, client, args.collection)
+        print("\nAnswer:\n")
+        print(answer)
+    else:
+        parser.print_help()
 
 if __name__ == "__main__":
     main()

src/retriever.py

@@ -0,0 +1,42 @@
+"""
+Retrieval and answer generation logic using LangChain.
+"""
+
+import os
+from typing import Any
+
+from chromadb import Client
+from chromadb.config import Settings
+
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.llms.openai import OpenAIChat
+from langchain.chains import RetrievalQA
+from langchain.vectorstores import Chroma
+
+def get_answer(question: str, client: Client, collection_name: str, k: int = 3) -> str:
+    """
+    Retrieve relevant FAQ chunks and generate an answer using OpenAIChat.
+    """
+    # Set up embeddings and LLM
+    embedding = OpenAIEmbeddings()
+    llm = OpenAIChat(temperature=0)
+
+    # Load vector store
+    vectorstore = Chroma(
+        client=client,
+        collection_name=collection_name,
+        embedding_function=embedding
+    )
+
+    # Build RetrievalQA chain
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=vectorstore.as_retriever(search_kwargs={"k": k}),
+        return_source_documents=True
+    )
+
+    # Run chain
+    result = qa_chain({"question": question})
+    answer = result.get("answer", "")
+    return answer.strip()