42 lines
1.1 KiB
Python
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() |