add drugs db index

alembic/versions/2025_08_07_2030-058739dc7aa6_add_drug_index.py

@@ -0,0 +1,36 @@
+"""add drug index
+
+Revision ID: 058739dc7aa6
+Revises: c509dcf806d4
+Create Date: 2025-08-07 20:30:35.962671
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '058739dc7aa6'
+down_revision: Union[str, Sequence[str], None] = 'c509dcf806d4'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('drugs', schema=None) as batch_op:
+        batch_op.create_index('ix_drugs_name_dosage_dosage_unit', ['name', 'dosage', 'dosage_unit'], unique=True)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('drugs', schema=None) as batch_op:
+        batch_op.drop_index('ix_drugs_name_dosage_dosage_unit')
+
+    # ### end Alembic commands ###

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "asyncpg>=0.30.0",
     "alembic>=1.16.4",
     "psycopg2-binary>=2.9.10",
+    "tqdm>=4.67.1",
 ]
 
 [project.optional-dependencies]

src/cache/drug_cache.py

@@ -6,6 +6,7 @@ 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
+from sqlalchemy.exc import IntegrityError
 
 
 class DrugFull(BaseModel):
@@ -36,17 +37,17 @@ async def convert_drug_result(drug: DrugPriceResponse) -> list[DrugFull]:
     return [DrugFull.model_validate(r) for r in response_dict["result"]]
 
 
-async def get_drug(drug_name: str) -> list[Drug]:
+async def get_drug(drug_name: str) -> list[DrugFull]:
     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:
+            try:
                 session.add(Drug(
                     name=drug.name,
                     dosage=drug.dosage,
@@ -55,6 +56,9 @@ async def store_drug(drugs: list[DrugFull]):
                     description=drug.description
                 ))
                 session.commit()
+            except IntegrityError as e:
+                session.rollback()
+                pass
 
 
 async def fetch_drug_with_dosage(drug_name: str, dosage: float) -> DrugFull | None:

src/config.py

@@ -29,6 +29,8 @@ class Settings(BaseSettings):
     POSTGRES_HOST: str = "db"
     POSTGRES_PORT: int = 5432
 
+    DEBUG: bool = False
+
 
     @property
     def DATABASE_URL(self):    # noqa

src/database.py

@@ -1,4 +1,4 @@
-from sqlalchemy import BigInteger, Column, MetaData, String, Float, Text
+from sqlalchemy import BigInteger, Column, MetaData, String, Float, Text, Index
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy import create_engine
@@ -17,5 +17,14 @@ class Drug(Base):
     unit_price = Column(Float, nullable=False)
     description = Column(Text, nullable=True)
 
+Index(
+    "ix_drugs_name_dosage_dosage_unit",
+    Drug.name,
+    Drug.dosage,
+    Drug.dosage_unit,
+    unique=True
+)
+
+
 engine = create_engine(settings.DATABASE_URL)
 Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)

src/services/session_service.py

@@ -61,7 +61,7 @@ class SessionService:
                 agent_id = self.agent_id or settings.TALESTORM_AGENT_ID
             
             response = await client.post("/sessions/", params={"agent_id": agent_id})
-            print(response.request.url)
+            # print(response.request.url)
             session_data = response.json()
             return str(session_data["id"])
 

uv.lock

@@ -426,6 +426,7 @@ dependencies = [
     { name = "redis" },
     { name = "requests" },
     { name = "sqlalchemy", extra = ["asyncio"] },
+    { name = "tqdm" },
 ]
 
 [package.optional-dependencies]
@@ -461,6 +462,7 @@ requires-dist = [
     { name = "redis", specifier = ">=6.2.0" },
     { name = "requests", specifier = ">=2.31.0" },
     { name = "sqlalchemy", extras = ["asyncio"], specifier = ">=2.0.42" },
+    { name = "tqdm", specifier = ">=4.67.1" },
 ]
 provides-extras = ["dev"]
 
@@ -1178,6 +1180,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f7/1f/b876b1f83aef204198a42dc101613fefccb32258e5428b5f9259677864b4/starlette-0.47.2-py3-none-any.whl", hash = "sha256:c5847e96134e5c5371ee9fac6fdf1a67336d5815e09eb2a01fdb57a351ef915b", size = 72984, upload-time = "2025-07-20T17:31:56.738Z" },
 ]
 
+[[package]]
+name = "tqdm"
+version = "4.67.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737, upload-time = "2024-11-24T20:12:22.481Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" },
+]
+
 [[package]]
 name = "typer"
 version = "0.16.0"