Merge pull request #3 from VectorInstitute/add_support_collections

Support multiple collections, separate medcat NER service

.env.development

@@ -1,6 +1,6 @@
 FRONTEND_PORT=8092
 BACKEND_PORT=8001
-CLINICAL_LLM_SERVICE_PORT=8003
+NER_SERVICE_PORT=8002
 BACKEND_HOST=backend-dev
 NODE_ENV=development
 BUILD_ID=latest

backend/Dockerfile.dev

@@ -13,9 +13,11 @@ RUN poetry config virtualenvs.create false && \
 ARG BACKEND_PORT
 ARG FRONTEND_PORT
 ARG JWT_SECRET_KEY
+ARG NER_SERVICE_PORT
 ENV BACKEND_PORT=${BACKEND_PORT}
 ENV FRONTEND_PORT=${FRONTEND_PORT}
 ENV JWT_SECRET_KEY=${JWT_SECRET_KEY}
+ENV NER_SERVICE_PORT=${NER_SERVICE_PORT}
 
 # Print port values
 RUN echo "FRONTEND_PORT: ${FRONTEND_PORT}"

backend/api/notes/data.py

@@ -1,6 +1,6 @@
 """Data models for medical notes."""
 
-from typing import List
+from typing import Dict, List
 
 from pydantic import BaseModel, Field
 
@@ -22,34 +22,85 @@ class MedicalNote(BaseModel):
     """
 
     note_id: str = Field(..., description="Unique identifier for the note")
-    subject_id: int = Field(..., description="Subject (patient) identifier")
-    hadm_id: str = Field(..., description="Hospital admission identifier")
+    patient_id: int = Field(..., description="Patient identifier")
+    encounter_id: str = Field(
+        ..., description="Hospital admission/encounter identifier"
+    )
     text: str = Field(..., description="Content of the medical note")
 
 
+class MetaAnnotation(BaseModel):
+    """
+    Represents a meta-annotation in a medical note.
+
+    Attributes
+    ----------
+    value: str
+        The value of the meta-annotation.
+    confidence: float
+        The confidence of the meta-annotation.
+    name: str
+
+    """
+
+    value: str
+    confidence: float
+    name: str
+
+
 class Entity(BaseModel):
     """
     Represents an entity in a medical note.
 
     Attributes
     ----------
-    entity_group : str
-        The type of the entity.
-    word : str
-        The word that the entity represents.
-    start : int
+    pretty_name: str
+        The pretty name of the entity.
+    cui: str
+        The CUI of the entity.
+    type_ids: List[str]
+        The type IDs of the entity.
+    types: List[str]
+        The types of the entity.
+    source_value: str
+        The source value of the entity.
+    detected_name: str
+        The detected name of the entity.
+    acc: float
+        The accuracy of the entity.
+    context_similarity: float
+        The context similarity of the entity.
+    start: int
         The start index of the entity in the text.
-    end : int
+    end: int
         The end index of the entity in the text.
-    score : float
-        The score of the entity.
+    icd10: List[Dict[str, str]]
+        The ICD-10 codes of the entity.
+    ontologies: List[str]
+        The ontologies of the entity.
+    snomed: List[str]
+        The SNOMED codes of the entity.
+    id: int
+        The ID of the entity.
+    meta_anns: Dict[str, MetaAnnotation]
+        The meta-annotations of the entity.
     """
 
-    entity_group: str
-    word: str
+    pretty_name: str
+    cui: str
+    type_ids: List[str]
+    types: List[str]
+    source_value: str
+    detected_name: str
+    acc: float
+    context_similarity: float
     start: int
     end: int
-    score: float
+    icd10: List[Dict[str, str]]
+    ontologies: List[str]
+    snomed: List[str]
+    id: int
+    meta_anns: Dict[str, MetaAnnotation]
 
 
 class NERResponse(BaseModel):

backend/api/notes/db.py

@@ -16,7 +16,7 @@
 MONGO_HOST = os.getenv("MONGO_HOST", "mongodb")
 MONGO_PORT = os.getenv("MONGO_PORT", "27017")
 MONGO_URL = f"mongodb://{MONGO_USERNAME}:{MONGO_PASSWORD}@{MONGO_HOST}:{MONGO_PORT}"
-DB_NAME = "medical_db"
+DB_NAME = "clinical_notes"
 
 
 async def get_database() -> AsyncIOMotorDatabase[Any]:
@@ -37,8 +37,11 @@ async def get_database() -> AsyncIOMotorDatabase[Any]:
     try:
         # Check the connection
         await client.admin.command("ismaster")
-        logger.info("Successfully connected to the database")
-        return client[DB_NAME]
+        db = client[DB_NAME]
+        logger.info(f"Successfully connected to the database: {DB_NAME}")
+        collections = await db.list_collection_names()
+        logger.info(f"Available collections: {collections}")
+        return db
     except Exception as e:
         logger.error(f"Unable to connect to the database: {str(e)}")
         raise ConnectionError(f"Database connection failed: {str(e)}") from e
@@ -58,12 +61,14 @@ async def check_database_connection() -> None:
         await client.admin.command("ismaster")
         db = client[DB_NAME]
         collections = await db.list_collection_names()
-        if "medical_notes" in collections:
+        if "mimiciv_discharge_notes" in collections:
             logger.info(
-                f"Database connection check passed. Found 'medical_notes' collection in {DB_NAME}"
+                f"Database connection check passed. Found 'mimiciv_discharge_notes' collection in {DB_NAME}"
             )
         else:
-            logger.warning(f"'medical_notes' collection not found in {DB_NAME}")
+            logger.warning(
+                f"'mimiciv_discharge_notes' collection not found in {DB_NAME}"
+            )
         logger.info("Database connection check passed")
     except Exception as e:
         logger.error(f"Database connection check failed: {str(e)}")