Create agent.py

alz_companion/agent.py

@@ -0,0 +1,333 @@
+from __future__ import annotations
+import os
+import json
+import base64
+import time
+import tempfile
+from typing import List, Dict, Any, Optional
+import re
+
+# OpenAI for LLM (optional)
+try:
+    from openai import OpenAI
+except Exception:
+    OpenAI = None
+
+# LangChain & RAG
+from langchain.schema import Document
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
+# TTS
+try:
+    from gtts import gTTS
+except Exception:
+    gTTS = None
+
+from .prompts import (
+    SYSTEM_TEMPLATE, ANSWER_TEMPLATE_CALM, ANSWER_TEMPLATE_ADQ,
+    SAFETY_GUARDRAILS, RISK_FOOTER, render_emotion_guidelines,
+    GOAL_ROUTER_PROMPT, SPECIALIST_CLASSIFICATION_PROMPT,
+    ANSWER_TEMPLATE_FACTUAL, ANSWER_TEMPLATE_GENERAL_KNOWLEDGE, ANSWER_TEMPLATE_GENERAL,
+    ROUTER_PROMPT as RAG_ROUTER_PROMPT
+)
+
+
+# -----------------------------
+# Multimodal Processing Functions
+# -----------------------------
+
+def _openai_client() -> Optional[OpenAI]:
+    api_key = os.getenv("OPENAI_API_KEY", "").strip()
+    return OpenAI(api_key=api_key) if api_key and OpenAI else None
+
+def describe_image(image_path: str) -> str:
+    """Uses a vision model to describe an image for context."""
+    client = _openai_client()
+    if not client:
+        return "(Image description failed: OpenAI API key not configured.)"
+    
+    try:
+        extension = os.path.splitext(image_path)[1].lower()
+        mime_type = "image/jpeg"
+        if extension == ".png": mime_type = "image/png"
+        elif extension in [".jpg", ".jpeg"]: mime_type = "image/jpeg"
+
+        with open(image_path, "rb") as image_file:
+            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
+
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[{"role": "user", "content": [
+                {"type": "text", "text": "Describe this image concisely for a memory journal. Focus on people, places, and key objects."},
+                {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}}
+            ]}],
+            max_tokens=100,
+        )
+        return response.choices[0].message.content or "No description available."
+    except Exception as e:
+        return f"[Image description error: {e}]"
+
+# -----------------------------
+# NLU Classification Function (Router/Specialist Model)
+# -----------------------------
+def detect_tags_from_query(
+    query: str,
+    behavior_options: list, emotion_options: list, topic_options: list, context_options: list,
+    example_retriever: Optional[FAISS] = None, **kwargs
+) -> Dict[str, Any]:
+    """Uses a dynamic few-shot process to classify the user's query."""
+    
+    print("\n--- RUNNING NLU V4 (Dynamic Few-Shot) ---")
+    
+    examples_str = "No examples provided."
+    if example_retriever:
+        try:
+            similar_docs = example_retriever.invoke(query)
+            examples = [doc.metadata["full_fixture"] for doc in similar_docs]
+            examples_str = "\n\n".join([
+                f"User Query: \"{ex['turns'][0]['text']}\"\n<thinking>\n{ex['expected'].get('reasoning', 'No reasoning provided.')}\n</thinking>\nJSON Response:\n{json.dumps(ex['expected'], indent=4)}"
+                for ex in examples
+            ])
+            print(f"Dynamically retrieved {len(examples)} examples for the prompt.")
+        except Exception as e:
+            print(f"Could not retrieve examples: {e}")
+
+    specialist_prompt = SPECIALIST_CLASSIFICATION_PROMPT.format(
+        behavior_options=", ".join(f'"{opt}"' for opt in behavior_options),
+        emotion_options=", ".join(f'"{opt}"' for opt in emotion_options),
+        topic_options=", ".join(f'"{opt}"' for opt in topic_options),
+        context_options=", ".join(f'"{opt}"' for opt in context_options),
+        examples=examples_str,
+        query=query
+    )
+    
+    messages = [{"role": "user", "content": specialist_prompt}]
+    response_str = call_llm(messages, temperature=0.1)
+    
+    print(f"\n--- NLU Full Response ---\n{response_str}\n-----------------------\n")
+    
+    result_dict = {"detected_behaviors": [], "detected_emotion": "None", "detected_topic": "None", "detected_contexts": []}
+    
+    try:
+        start_brace = response_str.find('{')
+        end_brace = response_str.rfind('}')
+        if start_brace != -1 and end_brace != -1 and end_brace > start_brace:
+            json_str = response_str[start_brace : end_brace + 1]
+            result = json.loads(json_str)
+            
+            result_dict["detected_behaviors"] = [b for b in result.get("detected_behaviors", []) if b in behavior_options]
+            result_dict["detected_emotion"] = result.get("detected_emotion") if result.get("detected_emotion") in emotion_options else "None"
+            result_dict["detected_topic"] = result.get("detected_topic") if result.get("detected_topic") in topic_options else "None"
+            result_dict["detected_contexts"] = [c for c in result.get("detected_contexts", []) if c in context_options]
+        
+        return result_dict
+    except (json.JSONDecodeError, AttributeError) as e:
+        print(f"ERROR parsing Specialist JSON: {e}")
+        return result_dict
+
+# -----------------------------
+# Embeddings & VectorStore
+# -----------------------------
+
+def _default_embeddings():
+    model_name = os.getenv("EMBEDDINGS_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+    return HuggingFaceEmbeddings(model_name=model_name)
+
+def build_or_load_vectorstore(docs: List[Document], index_path: str, is_personal: bool = False) -> FAISS:
+    os.makedirs(os.path.dirname(index_path), exist_ok=True)
+    if os.path.isdir(index_path) and os.path.exists(os.path.join(index_path, "index.faiss")):
+        try:
+            return FAISS.load_local(index_path, _default_embeddings(), allow_dangerous_deserialization=True)
+        except Exception as e:
+            print(f"Could not load existing vector store at {index_path}: {e}")
+    
+    if is_personal and not docs:
+        docs = [Document(page_content="(This is the start of the personal memory journal.)", metadata={"source": "placeholder"})]
+        
+    vs = FAISS.from_documents(docs, _default_embeddings())
+    vs.save_local(index_path)
+    return vs
+
+def texts_from_jsonl(path: str) -> List[Document]:
+    out: List[Document] = []
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            for i, line in enumerate(f):
+                line = line.strip()
+                if not line: continue
+                obj = json.loads(line)
+                txt = obj.get("text") or ""
+                if not isinstance(txt, str) or not txt.strip(): continue
+                md = {"source": os.path.basename(path), "chunk": i}
+                for k in ("behaviors", "emotion", "topic_tags", "context_tags"):
+                    if k in obj: md[k] = obj[k]
+                out.append(Document(page_content=txt, metadata=md))
+    except Exception as e:
+        print(f"Error reading from JSONL file {path}: {e}")
+        return []
+    return out
+
+def bootstrap_vectorstore(sample_paths: List[str] | None = None, index_path: str = "data/faiss_index") -> FAISS:
+    docs: List[Document] = []
+    for p in (sample_paths or []):
+        try:
+            if p.lower().endswith(".jsonl"):
+                docs.extend(texts_from_jsonl(p))
+            else:
+                with open(p, "r", encoding="utf-8", errors="ignore") as fh:
+                    docs.append(Document(page_content=fh.read(), metadata={"source": os.path.basename(p)}))
+        except Exception:
+            continue
+    if not docs:
+        docs = [Document(page_content="(empty index)", metadata={"source": "placeholder"})]
+    return build_or_load_vectorstore(docs, index_path=index_path)
+
+# -----------------------------
+# LLM Call
+# -----------------------------
+def call_llm(messages: List[Dict[str, str]], temperature: float = 0.6, stop: Optional[List[str]] = None) -> str:
+    client = _openai_client()
+    model = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+    if not client:
+        return "(Offline Mode: OpenAI API key not configured.)"
+    try:
+        api_args = {
+            "model": model, "messages": messages,
+            "temperature": float(temperature if temperature is not None else 0.6)
+        }
+        if stop: api_args["stop"] = stop
+        resp = client.chat.completions.create(**api_args)
+        return (resp.choices[0].message.content or "").strip()
+    except Exception as e:
+        return f"[LLM API Error: {e}]"
+
+# -----------------------------
+# Prompting & RAG Chain
+# -----------------------------
+
+def _format_sources(docs: List[Document]) -> List[str]:
+    return list(set(d.metadata.get("source", "unknown") for d in docs))
+
+def make_rag_chain(
+    vs_general: FAISS,
+    vs_personal: FAISS,
+    *,
+    role: str = "patient",
+    temperature: float = 0.6,
+    language: str = "English",
+    patient_name: str = "the patient",
+    caregiver_name: str = "the caregiver",
+    tone: str = "warm",
+    **kwargs
+):
+    """Returns a callable that performs the complete, RAG process."""
+    
+    def _format_docs(docs: List[Document], default_msg: str) -> str:
+        if not docs: return default_msg
+        unique_docs = {doc.page_content: doc for doc in docs}.values()
+        return "\n".join([f"- {d.page_content.strip()}" for d in unique_docs])
+
+    def _answer_fn(query: str, chat_history: List[Dict[str, str]], scenario_tag: Optional[str] = None, emotion_tag: Optional[str] = None, topic_tag: Optional[str] = None) -> Dict[str, Any]:
+        
+        router_messages = [{"role": "user", "content": RAG_ROUTER_PROMPT.format(query=query)}]
+        query_type = call_llm(router_messages, temperature=0.0).strip().lower()
+
+        system_message = SYSTEM_TEMPLATE.format(tone=tone, language=language, patient_name=patient_name or "the patient", caregiver_name=caregiver_name or "the caregiver", guardrails=SAFETY_GUARDRAILS)
+        messages = [{"role": "system", "content": system_message}]
+        messages.extend(chat_history)
+
+        if "factual_question" in query_type:
+            retriever_personal = vs_personal.as_retriever(search_kwargs={"k": 2})
+            retriever_general = vs_general.as_retriever(search_kwargs={"k": 2})
+            all_docs = retriever_personal.invoke(query) + retriever_general.invoke(query)
+            context = _format_docs(all_docs, "(No relevant information found in the memory journal.)")
+            user_prompt = ANSWER_TEMPLATE_FACTUAL.format(context=context, query=query, language=language)
+            messages.append({"role": "user", "content": user_prompt})
+            answer = call_llm(messages, temperature=temperature)
+            return {"answer": answer, "sources": _format_sources(all_docs)}
+
+        elif "general_knowledge_question" in query_type:
+            user_prompt = ANSWER_TEMPLATE_GENERAL_KNOWLEDGE.format(query=query, language=language)
+            messages.append({"role": "user", "content": user_prompt})
+            answer = call_llm(messages, temperature=temperature)
+            return {"answer": answer, "sources": ["General Knowledge"]}
+
+        elif "general_conversation" in query_type:
+            user_prompt = ANSWER_TEMPLATE_GENERAL.format(query=query, language=language)
+            messages.append({"role": "user", "content": user_prompt})
+            answer = call_llm(messages, temperature=temperature)
+            return {"answer": answer, "sources": []}
+
+        else: # Default to caregiving scenario
+            search_filter = {}
+            if scenario_tag and isinstance(scenario_tag, list):
+                search_filter["behaviors"] = [s.lower() for s in scenario_tag]
+            
+            if emotion_tag and emotion_tag != "None": search_filter["emotion"] = emotion_tag.lower()
+            if topic_tag and topic_tag != "None": search_filter["topic_tags"] = topic_tag.lower()
+
+            personal_docs = vs_personal.similarity_search(query, k=3, filter=search_filter if search_filter else None)
+            general_docs = vs_general.similarity_search(query, k=3, filter=search_filter if search_filter else None)
+            
+            personal_context = _format_docs(personal_docs, "(No relevant personal memories found.)")
+            general_context = _format_docs(general_docs, "(No general guidance found.)")
+            
+            all_docs = personal_docs + general_docs
+            first_emotion = next((doc.metadata.get("emotion") for doc in all_docs if doc.metadata.get("emotion")), emotion_tag)
+            emotions_context = render_emotion_guidelines(first_emotion)
+
+            is_tagged_scenario = scenario_tag or emotion_tag or first_emotion
+            template = ANSWER_TEMPLATE_ADQ if is_tagged_scenario else ANSWER_TEMPLATE_CALM
+            
+            display_scenario_tag = (scenario_tag[0] if isinstance(scenario_tag, list) and scenario_tag else scenario_tag) or ""
+            
+            if template == ANSWER_TEMPLATE_ADQ:
+                user_prompt = template.format(general_context=general_context, personal_context=personal_context, query=query, scenario_tag=display_scenario_tag, emotions_context=emotions_context, role=role, language=language)
+            else:
+                combined_context = f"General Guidance:\n{general_context}\n\nPersonal Memories:\n{personal_context}"
+                user_prompt = template.format(context=combined_context, query=query, language=language)
+
+            messages.append({"role": "user", "content": user_prompt})
+            answer = call_llm(messages, temperature=temperature)
+            
+            high_risk_scenarios = ["exit_seeking", "wandering", "elopement"]
+            if display_scenario_tag and display_scenario_tag.lower() in high_risk_scenarios:
+                answer += f"\n\n---\n{RISK_FOOTER}"
+
+            return {"answer": answer, "sources": _format_sources(all_docs)}
+
+    return _answer_fn
+
+def answer_query(chain, question: str, **kwargs) -> Dict[str, Any]:
+    if not callable(chain): return {"answer": "[Error: RAG chain is not callable]", "sources": []}
+    try:
+        return chain(question, **kwargs)
+    except Exception as e:
+        print(f"ERROR in answer_query: {e}")
+        return {"answer": f"[Error executing chain: {e}]", "sources": []}
+
+# -----------------------------
+# TTS & Transcription
+# -----------------------------
+def synthesize_tts(text: str, lang: str = "en"):
+    if not text or gTTS is None: return None
+    try:
+        fd, path = tempfile.mkstemp(suffix=".mp3")
+        os.close(fd)
+        tts = gTTS(text=text, lang=(lang or "en"))
+        tts.save(path)
+        return path
+    except Exception:
+        return None
+
+def transcribe_audio(filepath: str, lang: str = "en"):
+    client = _openai_client()
+    if not client:
+        return "[Transcription failed: API key not configured]"
+    api_args = {"model": "whisper-1"}
+    if lang and lang != "auto":
+        api_args["language"] = lang
+    with open(filepath, "rb") as audio_file:
+        transcription = client.audio.transcriptions.create(file=audio_file, **api_args)
+    return transcription.text