update: model choosing / more parameter data

.gitignore

@@ -1,3 +1,3 @@
 .venv
 .env
-__pycache__
+**/__pycache__/

DockerFile

@@ -10,4 +10,4 @@ COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -4,4 +4,5 @@ emoji: 🐳
 colorFrom: purple
 colorTo: gray
 sdk: docker
+app_port: 7860
 ---

cli.py

@@ -1,21 +1,27 @@
 from llm.client import NurseCLI
 from llm.llm import VirtualNurseLLM
+from dotenv import load_dotenv
+import os
 
-# model: typhoon-v1.5x-70b-instruct
-# nurse_llm = VirtualNurseLLM(
-#     base_url="https://api.opentyphoon.ai/v1",
-#     model="typhoon-v1.5x-70b-instruct",
-#     api_key=os.getenv("TYPHOON_API_KEY")
-# )
-
-# model: OpenThaiGPT
+load_dotenv()
 
 if __name__ == "__main__":
-    nurse_llm = VirtualNurseLLM(
-    base_url="https://api.aieat.or.th/v1",
-    model=".",
-    api_key="dummy"
-    )
+    model_choice = input("Choose the model to use (1 for typhoon-v1.5x-70b-instruct, 2 for OpenThaiGPT): ")
+    if model_choice == "1":
+        nurse_llm = VirtualNurseLLM(
+            base_url="https://api.opentyphoon.ai/v1",
+            model="typhoon-v1.5x-70b-instruct",
+            api_key=os.environ.get("TYPHOON_CHAT_KEY")
+        )
+    elif model_choice == "2":
+        nurse_llm = VirtualNurseLLM(
+            base_url="https://api.aieat.or.th/v1",
+            model="OpenThaiGPT",
+            api_key="dummy"
+        )
+    else:
+        print("Invalid choice. Exiting.")
+        exit(1)
 
     cli = NurseCLI(nurse_llm)
     cli.start()

llm/llm.py

@@ -19,8 +19,11 @@ class VirtualNurseLLM:
         self.JSON_EXAMPLE = JSON_EXAMPLE
         self.ehr_data = {}
         self.chat_history = []
+        self.current_patient_response = None
+        self.current_context = None
         self.debug = False
         self.current_prompt = None
+        self.current_prompt_ehr = None
         self.current_question = None
 
     def create_prompt(self, task_type):
@@ -41,7 +44,7 @@ class VirtualNurseLLM:
     def gather_ehr(self, patient_response, max_retries=3):
         prompt = self.create_prompt("extract_ehr")
         messages = prompt.format_messages(ehr_data=self.ehr_data, patient_response=patient_response, example=self.JSON_EXAMPLE)
-        self.current_prompt = messages
+        self.current_prompt_ehr = messages[0].content
         response = self.client(messages=messages)
         if self.debug:
             pprint(f"gather ehr llm response: \n{response.content}\n")
@@ -52,10 +55,10 @@ class VirtualNurseLLM:
                 json_content = self.extract_json_content(response.content)
                 if self.debug:
                     pprint(f"JSON after dumps:\n{json_content}\n")
-                ehr_data = EHRModel.parse_raw(json_content)
+                ehr_data = EHRModel.model_validate_json(json_content)
 
                 # Update only missing parameters
-                for key, value in ehr_data.dict().items():
+                for key, value in ehr_data.model_dump().items():
                     if value not in [None, [], {}]:  # Checks for None and empty lists or dicts
                         print(f"Updating {key} with value {value}")
                         self.ehr_data[key] = value
@@ -79,7 +82,7 @@ class VirtualNurseLLM:
                         retry_prompt=retry_prompt,
                         json_problem=json_content
                     )
-                    self.current_prompt = messages
+                    self.current_prompt_ehr = messages[0].content
                     print(f"กำลังลองใหม่ด้วย prompt ที่ปรับแล้ว: {retry_prompt}")
                     response = self.client(messages=messages)
 
@@ -109,7 +112,7 @@ class VirtualNurseLLM:
                 messages = ChatPromptTemplate.from_messages([question_prompt, history_context])
                 messages = messages.format_messages(description=f'"{field}":"{description}"', context=context, patient_response=patient_response, field_descriptions=self.field_descriptions)
                 self.current_context = context
-                self.current_prompt = messages
+                self.current_prompt = messages[0].content
                 
                 # format print
                 # pprint(pformat(messages.messages[0].prompt.template, indent=4, width=80))
@@ -127,6 +130,7 @@ class VirtualNurseLLM:
         if patient_response:
             self.chat_history.append({"role": "user", "content": patient_response})
         question = self.get_question(patient_response)
+        self.current_patient_response = patient_response
         self.chat_history.append({"role": "assistant", "content": question})
         return question
 

app.py → main.py

@@ -5,23 +5,23 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
 import os
-import dotenv
-dotenv.load_dotenv()
+from dotenv import load_dotenv
+load_dotenv()
 
 # model: typhoon-v1.5x-70b-instruct
-# nurse_llm = VirtualNurseLLM(
-#     base_url="https://api.opentyphoon.ai/v1",
-#     model="typhoon-v1.5x-70b-instruct",
-#     api_key=os.getenv("TYPHOON_API_KEY")
-# )
-
-# model: OpenThaiGPT
 nurse_llm = VirtualNurseLLM(
-    base_url="https://api.aieat.or.th/v1",
-    model=".",
-    api_key="dummy"
+    base_url="https://api.opentyphoon.ai/v1",
+    model="typhoon-v1.5x-70b-instruct",
+    api_key=os.getenv("TYPHOON_CHAT_KEY")
 )
 
+# model: OpenThaiGPT
+# nurse_llm = VirtualNurseLLM(
+#     base_url="https://api.aieat.or.th/v1",
+#     model=".",
+#     api_key="dummy"
+# )
+
 app = FastAPI()
 
 app.add_middleware(
@@ -34,7 +34,22 @@ app.add_middleware(
 
 class UserInput(BaseModel):
     user_input: str
+    model_name: str = "typhoon-v1.5x-70b-instruct"
+
+class NurseResponse(BaseModel):
+    nurse_response: str
 
+class EHRData(BaseModel):
+    ehr_data: dict
+    current_context: str
+    current_prompt: str
+    current_prompt_ehr: str
+    current_patient_response: str
+    current_question: str
+
+class ChatHistory(BaseModel):
+    chat_history: list
+    
 @app.get("/", response_class=HTMLResponse)
 def read_index():
     return """
@@ -53,21 +68,24 @@ def read_index():
 
 @app.get("/history")
 def get_chat_history():
-    return {"chat_history": nurse_llm.chat_history}
+    return ChatHistory(chat_history = nurse_llm.chat_history)
 
-@app.get("/ehr")
+@app.get("/details")
 def get_ehr_data():
-    return {"ehr_data": nurse_llm.ehr_data}
-
-@app.get("/status")
-def get_status():
-    return {"current_prompt": nurse_llm.current_prompt}
+    return EHRData(
+        ehr_data=nurse_llm.ehr_data,
+        current_context=nurse_llm.current_context,
+        current_prompt=nurse_llm.current_prompt,
+        current_prompt_ehr=nurse_llm.current_prompt_ehr,
+        current_patient_response=nurse_llm.current_patient_response,
+        current_question=nurse_llm.current_question
+    )
 
-@app.post("/debug")
 def toggle_debug():
     nurse_llm.debug = not nurse_llm.debug
     return {"debug_mode": "on" if nurse_llm.debug else "off"}
 
+
 @app.post("/reset")
 def data_reset():
     nurse_llm.reset()
@@ -75,8 +93,17 @@ def data_reset():
 
 @app.post("/nurse_response")
 def nurse_response(user_input: UserInput):
+    """
+    Models: "typhoon-v1.5x-70b-instruct (default)", "openthaigpt"
+    """
+    if user_input.model_name == "typhoon-v1.5x-70b-instruct":
+        nurse_llm.model = "typhoon-v1.5x-70b-instruct"
+    elif user_input.model_name == "openthaigpt":
+        nurse_llm.model = "openthaigpt"
+    else:
+        return {"error": "Invalid model name"}
     response = nurse_llm.invoke(user_input.user_input)
-    return {"nurse_response": response}
+    return NurseResponse(nurse_response = response)
 
 if __name__ == "__main__":
     uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)