File size: 3,725 Bytes
22e29e9 408dcaa 22e29e9 | 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 | import os
from huggingface_hub import AsyncInferenceClient
from core import config
from core.globals import ml_models
from engine.vector_store import CRMVectorStore
def load_rag_models():
"""
Initializes Qdrant vector store and the Hugging Face Inference client.
"""
if not config.HF_TOKEN:
print("⚠️ HF_TOKEN not set. RAG models will not be loaded.")
return
if not config.QDRANT_URL or not config.QDRANT_API_KEY:
print("⚠️ QDRANT_URL or QDRANT_API_KEY not set. Qdrant will not be loaded.")
else:
try:
store = CRMVectorStore()
ml_models["vector_store"] = store
print("✅ Qdrant vector store loaded.")
except Exception as e:
print(f"❌ Failed to initialize Qdrant vector store: {e}")
# Initialize the Serverless Inference API client
try:
client = AsyncInferenceClient(model=config.RAG_LLM_MODEL, token=config.HF_TOKEN)
ml_models["llm_client"] = client
print(f"✅ Hugging Face Inference Client created for {config.RAG_LLM_MODEL}")
except Exception as e:
print(f"❌ Failed to initialize HF Client: {e}")
async def get_rag_response(question: str):
store: CRMVectorStore = ml_models.get("vector_store")
client: AsyncInferenceClient = ml_models.get("llm_client")
if not store or not client:
raise RuntimeError("RAG models not loaded.")
# Retrieve top 5 contexts
hits = store.search(question, top_k=5)
if hits:
best_score = hits[0]["score"]
confidence = round(best_score, 2)
# Combine retrieved text
context_text = "\n\n".join([hit["payload"].get("text", "") for hit in hits])
else:
confidence = 0.0
context_text = "No relevant context found."
# Construct the prompt for Llama 3
system_prompt = "You are a helpful AI assistant. Answer the user's question based strictly on the provided context."
user_prompt = f"""
<context>
{context_text}
</context>
Question: {question}
If the answer is not in the context, just say, "I am sorry, but I cannot find the answer in the provided documents."
Answer:
"""
# Format messages for the Inference API Chat Completion
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
try:
# Call the HF Inference API asynchronously
response = await client.chat_completion(
messages=messages,
max_tokens=512,
temperature=0.3
)
answer = response.choices[0].message.content.strip()
except Exception as e:
answer = f"Error generating response from LLM API: {str(e)}"
return {"answer": answer, "confidence_score": confidence}
async def get_summary_response(text: str):
"""
Generates a summary for the given text using the HF Inference API.
"""
client: AsyncInferenceClient = ml_models.get("llm_client")
if not client:
raise RuntimeError("LLM Client not loaded.")
system_prompt = "You are an expert at summarizing text concisely."
user_prompt = f"Please summarize the following text:\n\n{text}\n\nSummary:"
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
try:
response = await client.chat_completion(
messages=messages,
max_tokens=512,
temperature=0.2
)
summary = response.choices[0].message.content.strip()
except Exception as e:
summary = f"Could not generate summary due to API error: {str(e)}"
return {"summary": summary}
|