Files
povtornyy-ekzamen-faq-bot-c…/src/retriever.py
T

42 lines
1.1 KiB
Python

"""
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()