Upload agent.py

agent.py

@@ -0,0 +1,171 @@
+# agent.py
+import os
+from typing import Dict, List, Any, Optional
+from langgraph.graph import Graph
+from langchain.schema import BaseMessage, HumanMessage, AIMessage
+from langchain_community.llms import HuggingFacePipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+import json
+
+
+class MedicalAgent:
+    def __init__(self):
+        self.llm = self._load_huatuogpt()
+        self.rag = MedicalRAG()
+        self.conversation_history = []
+        self.question_count = 0
+        self.max_questions = 3
+        self.max_words_per_question = 5
+
+    def _load_huatuogpt(self):
+        """Load HuatuoGPT model from HuggingFace"""
+        model_name = "HuatuoGPT/HuatuoGPT-7B"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=100,
+            temperature=0.7,
+            do_sample=True
+        )
+
+        return HuggingFacePipeline(pipeline=pipe)
+
+    def process_patient_input(self, patient_text: str) -> Dict[str, Any]:
+        """Process patient input and generate response"""
+        self.conversation_history.append(f"Patient: {patient_text}")
+
+        # Check if we've reached question limit
+        if self.question_count >= self.max_questions:
+            return self._generate_summary()
+
+        # Analyze symptoms and decide next action
+        analysis = self._analyze_symptoms(patient_text)
+
+        if analysis["needs_follow_up"]:
+            follow_up_question = self._generate_follow_up_question(analysis)
+            self.question_count += 1
+            return {
+                "type": "question",
+                "content": follow_up_question,
+                "question_count": self.question_count
+            }
+        else:
+            return self._generate_summary()
+
+    def _analyze_symptoms(self, patient_text: str) -> Dict[str, Any]:
+        """Analyze symptoms using RAG and LLM"""
+        # Search medical knowledge base
+        relevant_info = self.rag.search(patient_text, k=3)
+
+        prompt = f"""
+        Patient complaint: {patient_text}
+        Relevant medical information: {relevant_info}
+
+        Analyze the symptoms and determine:
+        1. If we need follow-up questions (True/False)
+        2. What key information is missing
+        3. Suggested follow-up questions (max 5 words each)
+
+        Respond in JSON format:
+        {{
+            "needs_follow_up": boolean,
+            "missing_info": list,
+            "possible_questions": list
+        }}
+        """
+
+        response = self.llm(prompt)
+        try:
+            analysis = json.loads(response)
+        except:
+            analysis = {
+                "needs_follow_up": True,
+                "missing_info": ["symptom details"],
+                "possible_questions": ["How long have headache?", "Any other symptoms?"]
+            }
+
+        return analysis
+
+    def _generate_follow_up_question(self, analysis: Dict) -> str:
+        """Generate concise follow-up question"""
+        possible_questions = analysis.get("possible_questions", [])
+        if possible_questions:
+            question = possible_questions[0]
+            # Ensure question is within word limit
+            words = question.split()[:self.max_words_per_question]
+            return " ".join(words)
+        else:
+            return "Any other symptoms?"
+
+    def _generate_summary(self) -> Dict[str, Any]:
+        """Generate summary for doctor"""
+        conversation_text = "\n".join(self.conversation_history)
+
+        prompt = f"""
+        Patient conversation:
+        {conversation_text}
+
+        Generate a concise medical summary for the doctor including:
+        - Main symptoms
+        - Key findings
+        - Suggested preliminary diagnosis
+        - Recommended tests if any
+
+        Keep it under 150 words.
+        """
+
+        summary = self.llm(prompt)
+        return {
+            "type": "summary",
+            "content": summary,
+            "question_count": self.question_count
+        }
+
+    def process_doctor_question(self, doctor_text: str) -> str:
+        """Process doctor's follow-up questions"""
+        prompt = f"""
+        Doctor's question: {doctor_text}
+
+        Rephrase this question to be clear and simple for the patient.
+        Keep it under 5 words and make it easy to understand.
+        """
+
+        simplified_question = self.llm(prompt)
+        return simplified_question
+
+
+class MedicalRAG:
+    def __init__(self):
+        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
+        self.index = faiss.IndexFlatL2(384)
+        self.knowledge_base = []
+
+    def add_medical_knowledge(self, documents: List[str]):
+        """Add medical knowledge documents to RAG"""
+        self.knowledge_base.extend(documents)
+        embeddings = self.encoder.encode(documents)
+        self.index.add(np.array(embeddings))
+
+    def search(self, query: str, k: int = 3) -> List[str]:
+        """Search medical knowledge base"""
+        query_embedding = self.encoder.encode([query])
+        distances, indices = self.index.search(query_embedding, k)
+
+        results = []
+        for idx in indices[0]:
+            if idx < len(self.knowledge_base):
+                results.append(self.knowledge_base[idx])
+
+        return results