Spaces:
Sleeping
Sleeping
File size: 4,468 Bytes
835ecb4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate
from typing import Optional
import os
class RAGChain:
"""RAG chain using OpenAI API with Russian language support."""
def __init__(self,
retriever,
model_name: str = "gpt-4o-mini",
temperature: float = 0.3,
api_key: Optional[str] = None):
"""
Initialize RAG chain.
Args:
retriever: LangChain retriever (from vector store)
model_name: OpenAI model name
temperature: Temperature for LLM
api_key: OpenAI API key
"""
self.llm = ChatOpenAI(
model_name=model_name,
temperature=temperature,
api_key=api_key or os.getenv("OPENAI_API_KEY"),
max_tokens=1024
)
self.retriever = retriever
# Custom prompt for Russian language
self.prompt_template = PromptTemplate(
template="""Вы - полезный ассистент, специализирующийся на анализе документов.
Используя следующий контекст из документов, ответьте на вопрос.
Контекст:
{context}
Вопрос: {question}
Инструкции:
1. Ответьте только на основе информации из контекста
2. Если информация не найдена в контексте, скажите "Информация не найдена в документах"
3. Ответьте на русском языке
4. Будьте кратким и точным
5. Цитируйте источники если возможно
Ответ:""",
input_variables=["context", "question"]
)
# Create RetrievalQA chain
self.chain = RetrievalQA.from_chain_type(
llm=self.llm,
chain_type="stuff",
retriever=self.retriever,
return_source_documents=True,
chain_type_kwargs={"prompt": self.prompt_template}
)
def query(self, question: str) -> dict:
"""
Query the RAG chain.
Args:
question: User question (can be in any language)
Returns:
Dictionary with answer and source documents
"""
try:
result = self.chain.invoke({"query": question})
return {
"answer": result.get("result", ""),
"sources": [
{
"content": doc.page_content[:200], # First 200 chars
"metadata": doc.metadata
}
for doc in result.get("source_documents", [])
]
}
except Exception as e:
return {
"answer": f"Ошибка при обработке запроса: {str(e)}",
"sources": []
}
def query_with_context(self, question: str, context_limit: int = 5) -> dict:
"""
Query with explicit context retrieval.
Args:
question: User question
context_limit: Number of context chunks to retrieve
Returns:
Dictionary with answer and context
"""
# Retrieve relevant documents
relevant_docs = self.retriever.get_relevant_documents(
question,
search_kwargs={"k": context_limit}
)
# Format context
context = "\n\n".join([
f"Источник: {doc.metadata}\n{doc.page_content}"
for doc in relevant_docs
])
# Create prompt
prompt = self.prompt_template.format(context=context, question=question)
# Get response
response = self.llm.invoke(prompt)
return {
"answer": response.content,
"context_documents": [
{
"content": doc.page_content[:300],
"metadata": doc.metadata
}
for doc in relevant_docs
]
} |