add: MedQAAssistant

medrag_multi_modal/assistant/__init__.py

@@ -1,3 +1,4 @@
 from .llm_client import LLMClient
+from .medqa_assistant import MedQAAssistant
 
-__all__ = ["LLMClient"]
+__all__ = ["LLMClient", "MedQAAssistant"]

medrag_multi_modal/assistant/llm_client.py

@@ -29,7 +29,7 @@ class LLMClient(weave.Model):
         schema: Optional[Any] = None,
     ) -> Union[str, Any]:
         import google.generativeai as genai
-        
+
         system_prompt = (
             [system_prompt] if isinstance(system_prompt, str) else system_prompt
         )

medrag_multi_modal/assistant/medqa_assistant.py

@@ -0,0 +1,28 @@
+from typing import Optional
+
+import weave
+from PIL import Image
+
+from ..retrieval import SimilarityMetric
+from .llm_client import LLMClient
+
+
+class MedQAAssistant(weave.Model):
+    llm_client: LLMClient
+    retriever: weave.Model
+    top_k_chunks: int = 2
+    retrieval_similarity_metric: SimilarityMetric = SimilarityMetric.COSINE
+
+    @weave.op()
+    def predict(self, query: str, image: Optional[Image.Image] = None) -> str:
+        _image = image
+        retrieved_chunks = self.retriever.predict(
+            query, top_k=self.top_k_chunks, metric=self.retrieval_similarity_metric
+        )
+        retrieved_chunks = [chunk["text"] for chunk in retrieved_chunks]
+        system_prompt = """
+        You are a medical expert. You are given a query and a list of chunks from a medical document.
+        """
+        return self.llm_client.predict(
+            system_prompt=system_prompt, user_prompt=retrieved_chunks
+        )