File size: 1,817 Bytes
2303ac9 |
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 |
import sys
from pathlib import Path
from typing import List, Dict
import litellm
# Add project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
import config
class LLMHandler:
def __init__(self, model_name: str = None):
self.model_name = model_name if model_name else config.LLM_MODEL_NAME
def generate_answer(self, query: str, context_chunks: List[Dict], api_key: str, temperature: float = 0.3) -> str:
"""
Generates an answer using the LLM and the provided context.
"""
if not context_chunks:
return "Не знайдено релевантних документів для відповіді."
# Prepare context string
context_text = ""
for i, chunk in enumerate(context_chunks):
meta = chunk['chunk']['metadata']
cit = meta.get('citation_short', f"Doc {meta.get('article_number')}")
title = meta.get('article_title', '')
text = chunk['chunk']['text']
context_text += f"Source [{i+1}]: {cit}"
if title:
context_text += f" ({title})"
context_text += f"\nContent: {text}\n\n"
# Construct messages
messages = [
{"role": "system", "content": config.SYSTEM_PROMPT.format(context=context_text)},
{"role": "user", "content": query}
]
try:
response = litellm.completion(
model=self.model_name,
messages=messages,
api_key=api_key,
temperature=temperature,
max_tokens=1024
)
return response.choices[0].message.content
except Exception as e:
return f"Error communicating with LLM: {str(e)}"
|