""" Module for RAG (Retrieval Augmented Generation) functionality. """ import sys import os sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from typing import List import logging import os import inspect from openai import OpenAI from pydantic import BaseModel from dotenv import load_dotenv from configs.config import external_client, model_config load_dotenv() class QueryRequest(BaseModel): query: str class QueryResponse(BaseModel): answer: str sources: List[str] = [] metadata: dict = {} class SelectionRequest(BaseModel): selected_text: str question: str class SelectionResponse(BaseModel): answer: str class RAGService: def __init__(self, vector_store_collection: str = "Humanoids", vector_store_timeout: int = 30): """ Initialize the RAG service with required components. Args: vector_store_collection: Name of the vector store collection to use vector_store_timeout: Timeout for vector store operations """ try: # Use the centralized client configuration self.client = external_client self.model = model_config.model if hasattr(model_config, 'model') else 'qwen3-coder-plus' # Initialize vector store with proper configuration from data.vector_store import VectorStore self.vector_store = VectorStore( collection_name=vector_store_collection, timeout=vector_store_timeout ) # Check vector store connection on initialization # Note: check_connection is synchronous, so we can call it during init if not self.vector_store.check_connection(): logging.warning("Could not establish connection to vector store during initialization") except Exception as e: logging.error(f"Failed to initialize RAG service: {e}") raise async def check_api_connection(self): """ Check if the Qwen API connection is working. Returns: True if connection is successful, False otherwise """ try: # Check if this is an AsyncOpenAI client (which has async create method) # For AsyncOpenAI clients, the create method itself is async response = await self.client.chat.completions.create( model=self.model, messages=[ {"role": "system", "content": "You are a test assistant."}, {"role": "user", "content": "Test connection"} ], max_tokens=10 ) return True except Exception as e: logging.error(f"API connection test failed: {e}") return False def check_vector_store_connection(self): """ Check if the vector store connection is working. Returns: True if connection is successful, False otherwise """ try: return self.vector_store.check_connection() if self.vector_store else False except Exception as e: logging.error(f"Vector store connection test failed: {e}") return False async def check_health(self): """ Check the health of both the API and vector store. Returns: Dictionary with health status of both components """ try: api_healthy = await self.check_api_connection() vector_store_healthy = self.check_vector_store_connection() # Get vector store health details if available vector_store_details = {} if self.vector_store: try: vector_store_details = self.vector_store.health_check() except Exception: vector_store_details = {"status": "unreachable"} return { "api_healthy": api_healthy, "vector_store_healthy": vector_store_healthy, "vector_store_details": vector_store_details, "status": "healthy" if api_healthy and vector_store_healthy else "unhealthy" } except Exception as e: logging.error(f"Health check failed: {e}") return { "api_healthy": False, "vector_store_healthy": False, "status": "unhealthy", "error": str(e) } def set_vector_store(self, vector_store): """ Set the vector store instance to use for retrieval. Args: vector_store: Initialized VectorStore instance """ self.vector_store = vector_store async def query(self, query_text: str, top_k: int = 5) -> QueryResponse: """ Process a query using the RAG pipeline. Args: query_text: User's question top_k: Number of top results to retrieve from vector store Returns: QueryResponse with answer and sources """ try: # Validate inputs if not query_text or not query_text.strip(): return QueryResponse(answer="Please provide a query", sources=[], metadata={}) # Step 1: Generate embedding for the query from data.embeddings import EmbeddingService embedding_service = EmbeddingService() query_embedding = embedding_service.embed_text(query_text) # Step 2: Verify vector store is available and retrieve relevant documents if not self.vector_store: raise ValueError("Vector store not set in RAG service") # Check connection before performing search if not self.vector_store.check_connection(): logging.error("Vector store connection lost during query") return QueryResponse(answer="Service temporarily unavailable", sources=[], metadata={}) retrieved_docs = self.vector_store.search(query_embedding, limit=top_k) if not retrieved_docs: return QueryResponse(answer="I don't know", sources=[], metadata={"retrieved_docs_count": 0}) # Step 3: Format context from retrieved documents context = "\n\n".join([doc["content"] for doc in retrieved_docs]) sources = list(set([doc["doc_path"] for doc in retrieved_docs])) # Unique sources # Step 4: Generate answer using OpenAI with the context prompt = f""" Context information is below: {context} Using the provided context information, answer the question: {query_text} If the context does not contain sufficient information to answer the question, respond with "I don't know". """ # For AsyncOpenAI clients (like Qwen), the create method itself is async response = await self.client.chat.completions.create( model=self.model, messages=[ {"role": "system", "content": "You are a helpful assistant that answers questions based only on the provided context. If the context does not contain sufficient information to answer the question, respond with 'I don't know'."}, {"role": "user", "content": prompt} ], max_tokens=500, temperature=0.1 ) answer = response.choices[0].message.content.strip() # Check if the answer is "I don't know" if "i don't know" in answer.lower(): return QueryResponse(answer="I don't know", sources=[], metadata={"retrieved_docs_count": len(retrieved_docs)}) return QueryResponse( answer=answer, sources=sources, metadata={ "retrieved_docs_count": len(retrieved_docs), "sources_count": len(sources) } ) except Exception as e: logging.error(f"Error in query processing: {e}") return QueryResponse( answer="An error occurred while processing your query", sources=[], metadata={"error": str(e)} ) async def answer_from_selection(self, selected_text: str, question: str) -> SelectionResponse: """ Answer a question based only on the selected text. Args: selected_text: Text selected by the user question: Question about the selected text Returns: SelectionResponse with answer """ try: # Validate inputs if not selected_text or not selected_text.strip(): return SelectionResponse(answer="Please provide selected text") if not question or not question.strip(): return SelectionResponse(answer="Please provide a question") # Format the prompt using only the selected text prompt = f""" Selected text is below: {selected_text} Using only the provided selected text, answer the question: {question} If the selected text does not contain sufficient information to answer the question, respond with "I don't know". """ # For AsyncOpenAI clients (like Qwen), the create method itself is async response = await self.client.chat.completions.create( model=self.model, messages=[ {"role": "system", "content": "You are a helpful assistant that answers questions based only on the provided selected text. If the selected text does not contain sufficient information to answer the question, respond with 'I don't know'."}, {"role": "user", "content": prompt} ], max_tokens=500, temperature=0.1 ) answer = response.choices[0].message.content.strip() # Check if the answer is "I don't know" if "i don't know" in answer.lower(): return SelectionResponse(answer="I don't know") return SelectionResponse(answer=answer) except Exception as e: logging.error(f"Error in selection-based answering: {e}") return SelectionResponse(answer="An error occurred while processing your question")