add drug dosages parsing, add postgres db

src/cache/drug_cache.py

@@ -0,0 +1,90 @@
+
+import json
+import httpx
+from pydantic import BaseModel
+from src.database import Drug, Session
+from src.drug_price_parser import DrugPriceParser, DrugPriceResponse
+from src.config import settings
+from src.services.session_service import session_service
+
+
+class DrugFull(BaseModel):
+    name: str
+    dosage: float
+    dosage_unit: str
+    unit_price: float
+    description: str | None = None
+
+
+async def convert_drug_result(drug: DrugPriceResponse) -> list[DrugFull]:
+    base_url = settings.TALESTORM_API_BASE_URL
+    api_key = settings.TALESTORM_API_KEY
+
+    client = httpx.AsyncClient(
+        base_url=base_url,
+        headers={"X-API-Key": api_key},
+        timeout=httpx.Timeout(60.0, connect=10.0)  # 30s total timeout, 10s connect timeout
+    )
+    session_id = await session_service.create_session(agent_id=settings.TALESTORM_DRUG_AGENT_ID)
+    drug_json = drug.model_dump_json()
+    response = await client.post(
+        "/chat/",
+        json={"chat_session_id": session_id, "user_message": f"{drug_json}"},
+    )
+    response_json = response.json()['message']
+    response_dict = json.loads(response_json)
+    return [DrugFull.model_validate(r) for r in response_dict["result"]]
+
+
+async def get_drug(drug_name: str) -> list[Drug]:
+    parser = DrugPriceParser()
+    result = parser.get_drug_prices(drug_name)
+    drugs = await convert_drug_result(result)
+    
+    return drugs
+
+
+async def store_drug(drugs: list[DrugFull]):
+    with Session() as session:
+        for drug in drugs:
+            session.add(Drug(
+                name=drug.name,
+                dosage=drug.dosage,
+                dosage_unit=drug.dosage_unit,
+                unit_price=drug.unit_price,
+                description=drug.description
+            ))
+        session.commit()
+
+
+async def fetch_drug_with_dosage(drug_name: str, dosage: float) -> DrugFull | None:
+    try:
+        with Session() as session:
+            drug = session.query(Drug).filter(Drug.name == drug_name, Drug.dosage == dosage).first()
+            if drug:
+                return DrugFull.model_validate(drug)
+            
+            drug = session.query(Drug).filter(Drug.name == drug_name).first()
+            if drug:
+                return DrugFull.model_validate(drug)
+    except:
+        pass
+
+    drugs = await get_drug(drug_name)
+    print(f"Drug {drug_name} found {drugs}")
+    try:
+        await store_drug(drugs)
+    except Exception as e:
+        print(f"Error storing drug {drug_name}: {e}")
+        pass
+    
+    drug = None
+    for c_drug in drugs:
+        if c_drug.dosage == dosage:
+            drug = c_drug
+
+    if drug:
+        return DrugFull.model_validate(drug)
+
+    raise Exception(f"Drug {drug_name} with dosage {dosage} not found")
+

src/config.py

@@ -14,13 +14,26 @@ class Settings(BaseSettings):
     TALESTORM_AGENT_ID: str
     TALESTORM_DTQ_AGENT_ID: str
     TALESTORM_ESTIMATION_AGENT_ID: str
+    TALESTORM_DRUG_AGENT_ID: str
     VIRGIL_API_BASE_URL: str
+
     REDIS_HOST: str = "redis"
     REDIS_PORT: int = 6379
     REDIS_DB: int = 0
     REDIS_CACHE_TTL: int = 3600 * 24
     REDIS_CACHE_TTL_DRUGS: int = 3600 * 24 * 30
 
+    POSTGRES_USER: str
+    POSTGRES_PASSWORD: str
+    POSTGRES_DB: str
+    POSTGRES_HOST: str = "db"
+    POSTGRES_PORT: int = 5432
+
+
+    @property
+    def DATABASE_URL(self):    # noqa
+        return f'postgresql://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}'
+
 
 # Global settings instance
 settings = Settings(_env_file=env_path, _env_file_encoding='utf-8') 

src/database.py

@@ -0,0 +1,21 @@
+from sqlalchemy import BigInteger, Column, MetaData, String, Float, Text
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import create_engine
+from src.config import settings
+
+class Base(DeclarativeBase):
+    metadata = MetaData()
+
+class Drug(Base):
+    __tablename__ = "drugs"
+    
+    id = Column(BigInteger, primary_key=True, index=True)
+    name = Column(String, nullable=False, index=True)
+    dosage = Column(Float, nullable=False)
+    dosage_unit = Column(String, nullable=False)
+    unit_price = Column(Float, nullable=False)
+    description = Column(Text, nullable=True)
+
+engine = create_engine(settings.DATABASE_URL)
+Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)

src/drug_price_parser.py

@@ -87,13 +87,18 @@ class DrugPriceParser:
             dosages_table = formulation.find_next('div')
             dosage_elements = dosages_table.find_all('details')
             for dosage in dosage_elements:
+                quantity_table = dosage.find('table')
+                quantity_row = quantity_table.find_all('td', {'class': 'ddc-text-right'})
+                price_per_unit = quantity_row[0].get_text()
+                price_per_unit = price_per_unit.split(" ")[0]
+
                 summary = dosage.find('summary')
                 spans = summary.find_all('span')
                 dosage_name = spans[0].find('b').get_text()
-                dosage_price = spans[1].find_next('b').get_text()
+                # dosage_price = spans[1].find_next('b').get_text()
                 dosages.append(Dosage(
                     name=dosage_name.rstrip(),
-                    price=float(dosage_price.rstrip().replace('$', '').replace(',', ''))
+                    price=float(price_per_unit.rstrip().replace('$', '').replace(',', ''))
                 ))
             formulations.append(Formulation(
                 name=formulation_name.rstrip(),

src/models.py

@@ -124,4 +124,19 @@ class HealthResponse(BaseModel):
     status: str
     service: str
     version: str
-    dependencies: dict 
+    dependencies: dict
+
+class DrugCreate(BaseModel):
+    name: str = Field(..., description="Drug name")
+    dosage: float = Field(..., description="Dosage amount")
+    dosage_unit: str = Field(..., description="Dosage unit (e.g., mg, ml)")
+    unit_price: float = Field(..., description="Price per unit")
+    description: Optional[str] = Field(None, description="Drug description")
+
+class DrugResponse(BaseModel):
+    id: int
+    name: str
+    dosage: float
+    dosage_unit: str
+    unit_price: float
+    description: Optional[str] 

src/services/estimation_service_v2.py

@@ -1,8 +1,9 @@
 from datetime import date
 from enum import Enum
 from typing import Optional
+from src.cache.drug_cache import fetch_drug_with_dosage
 from src.models import PHQ, Applicant, Plan, EstimationResponse, EstimationDetails, EstimationResult
-from src.cache.redis_cache import fetch_conditions, fetch_drug, get_plan_by_id, search_drug
+from src.cache.redis_cache import fetch_conditions, get_plan_by_id, search_drug
 from src.config import settings
 import httpx
 
@@ -159,7 +160,7 @@ class EstimationService:
         today = date.today()
         return today.year - born.year - ((today.month, today.day) < (born.month, born.day))
 
-    def calculate_rx_spend(self, phq: PHQ, applicant_id: int) -> float:
+    async def calculate_rx_spend(self, phq: PHQ, applicant_id: int) -> float:
         rx_spend = 0
         for medication in phq.medications:
             if medication.applicant != applicant_id:
@@ -167,9 +168,29 @@ class EstimationService:
             try:
                 drug_name = medication.name
                 drug_url = search_drug(drug_name)
-                drug_price = fetch_drug(drug_url)
-                rx_spend += drug_price.prices[0].dosages[0].price
-            except Exception:
+                drug_dosage = float(medication.dosage)
+                drug_price = await fetch_drug_with_dosage(drug_url, drug_dosage)
+
+                if medication.frequency in ["Once daily", "At bedtime"]:
+                    month_times = 30
+                elif medication.frequency == "Twice daily":
+                    month_times = 60
+                elif medication.frequency in ["Three times daily", "After meals", "Before meals"]:
+                    month_times = 90
+                elif medication.frequency == "Four times daily":
+                    month_times = 120
+                elif medication.frequency == "Weekly":
+                    month_times = 4
+                elif medication.frequency == "Monthly":
+                    month_times = 1
+                elif medication.frequency == "Every other day":
+                    month_times = 15
+                else:
+                    month_times = 1
+
+                rx_spend += drug_price.unit_price * month_times
+            except Exception as e:
+                raise e
                 pass
         return rx_spend
 
@@ -241,7 +262,7 @@ class EstimationService:
                     is_dtq = True
                     reason = "Declined due to high BMI of one or more applicants"
             
-            rx_spend_applicant = self.calculate_rx_spend(phq, applicant_id)
+            rx_spend_applicant = await self.calculate_rx_spend(phq, applicant_id)
             rx_spend += rx_spend_applicant
 
             applicant_new_tier = self.get_tier(plans[0].coverage, rx_spend_applicant)
@@ -262,7 +283,7 @@ class EstimationService:
                     bmi=self.calculate_bmi(applicant.weight, applicant.heightFt, applicant.heightIn),
                     tier=applicant_tier.value,
                     rx_spend=rx_spend_applicant,
-                    message=reason if reason else f"Tier {applicant_tier} assigned with Rx spend within allowed limits."
+                    message=reason if reason else f"Tier {applicant_tier.value} assigned with Rx spend within allowed limits."
                 )
             )