Jemayz/Atlast

1import logging
2import json
3import re
4from src.doc_qa import AgenticQA
5from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
6from langchain_google_genai import ChatGoogleGenerativeAI
7
8logger = logging.getLogger(__name__)
9
10def load_rag_system(collection_name,domain):
11    """
12    Loads an existing RAG system by connecting to the persistent vector store.
13    This is fast and does not re-process any documents.
14    """
15    logger.info(f"Loading RAG system for collection: '{collection_name}'  (Domain: {domain})...")
16    try:
17        agent = AgenticQA(
18            config={
19                "retriever": {
20                    "collection_name": collection_name,
21                    "persist_directory": "chroma_db"
22                },
23                "domain": domain
24            }
25        )
26        # Check if the agent was actually created
27        if not agent.agent_executor:
28            raise Exception("Agent Executor was not created. Check logs for errors.")
29
30        logger.info(f"✅ System for '{collection_name}' loaded successfully.")
31        return agent
32    except Exception as e:
33        logger.error(f"❌ Failed to load RAG system for '{collection_name}': {e}")
34        logger.warning("Did you run the ingest.py script first?")
35        return None
36
37def markdown_bold_to_html(text: str):
38    """Converts markdown bold syntax to HTML <strong> tags."""
39    return re.sub(r"\*\*(.*?)\*\*", r"<strong>\1</strong>", text)
40
41def standardize_query(query):
42    if not query:
43        return None
44    return query.strip().lower()
45
46def get_standalone_question(input_question, chat_history,llm):
47    """Uses LLM to create a standalone question from the chat history."""
48    if not chat_history:
49        return input_question
50    
51    contextualize_q_prompt = ChatPromptTemplate.from_messages([
52        ("system", "Given a chat history and the latest user question which might reference context in the chat history, "
53        "formulate a standalone question which can be understood without the chat history. "
54        "IMPORTANT: DO NOT PROVIDE ANY ANSWERS. ONLY REPHRASE THE QUESTION IF NEEDED. "
55        "If the question is already clear and standalone, return it exactly as is. "
56        "Output ONLY the reformulated question, nothing else."),
57        MessagesPlaceholder("chat_history"),
58        ("human", "{input}"),
59    ])
60    history_aware_retriever_chain = contextualize_q_prompt | llm
61    
62    response = history_aware_retriever_chain.invoke(
63        {"chat_history": chat_history, "input": input_question}
64    )
65    return response.content
66
67def parse_agent_response(response_dict):
68    """A robust helper to parse the dictionary from an AgenticQA agent."""
69    answer = markdown_bold_to_html(response_dict.get('answer', 'Error: No answer found.'))
70    thoughts = response_dict.get('thoughts', 'No thought process available.')
71    validation = response_dict.get('validation', (False, 'Validation failed.'))
72    source = response_dict.get('source', 'Unknown')
73    
74    if validation and validation[1] == "Validation skipped for insurance domain.":
75        validation = (True, "Factual Answer")
76        
77    return answer, thoughts,validation, source
78    
79def extract_json_from_string(text: str) -> dict:
80    """
81    Finds and parses the first valid JSON object within a string.
82    Returns a dictionary, or an empty dict if no JSON is found.
83    """
84    # This regex finds the first occurrence of a string starting with { and ending with }
85    json_match = re.search(r'\{.*\}', text, re.DOTALL)
86    
87    if json_match:
88        json_string = json_match.group(0)
89        try:
90            return json.loads(json_string)
91        except json.JSONDecodeError:
92            # The extracted string is not valid JSON
93            return {"error": "Failed to parse extracted JSON", "raw_text": json_string}
94    else:
95        # No JSON object found in the string
96        return {"error": "No JSON object found in the string", "raw_text": text}