adding proper cost estimation and initial data logs

requirements.txt

@@ -0,0 +1,373 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements.txt
+aiofiles==24.1.0
+    # via gradio
+aiohappyeyeballs==2.6.1
+    # via aiohttp
+aiohttp==3.13.2
+    # via langchain-community
+aiosignal==1.4.0
+    # via aiohttp
+annotated-doc==0.0.4
+    # via fastapi
+annotated-types==0.7.0
+    # via pydantic
+anthropic==0.75.0
+    # via langchain-anthropic
+anyio==4.12.0
+    # via
+    #   anthropic
+    #   gradio
+    #   httpx
+    #   openai
+    #   starlette
+attrs==25.4.0
+    # via aiohttp
+audioop-lts==0.2.2
+    # via gradio
+brotli==1.2.0
+    # via gradio
+cachetools==6.2.2
+    # via google-auth
+certifi==2025.11.12
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+charset-normalizer==3.4.4
+    # via requests
+click==8.3.1
+    # via
+    #   typer
+    #   typer-slim
+    #   uvicorn
+dataclasses-json==0.6.7
+    # via langchain-community
+distro==1.9.0
+    # via
+    #   anthropic
+    #   openai
+docstring-parser==0.17.0
+    # via anthropic
+fastapi==0.123.4
+    # via gradio
+ffmpy==1.0.0
+    # via gradio
+filelock==3.20.0
+    # via huggingface-hub
+filetype==1.2.0
+    # via langchain-google-genai
+frozenlist==1.8.0
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2025.10.0
+    # via
+    #   gradio-client
+    #   huggingface-hub
+google-ai-generativelanguage==0.9.0
+    # via langchain-google-genai
+google-api-core==2.28.1
+    # via google-ai-generativelanguage
+google-auth==2.43.0
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+googleapis-common-protos==1.72.0
+    # via
+    #   google-api-core
+    #   grpcio-status
+gradio==6.0.2
+    # via litigation-research (pyproject.toml)
+gradio-client==2.0.1
+    # via gradio
+greenlet==3.2.4
+    # via sqlalchemy
+groovy==0.1.2
+    # via gradio
+grpcio==1.76.0
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+    #   grpcio-status
+grpcio-status==1.76.0
+    # via google-api-core
+h11==0.16.0
+    # via
+    #   httpcore
+    #   uvicorn
+hf-xet==1.2.0
+    # via huggingface-hub
+httpcore==1.0.9
+    # via httpx
+httpx==0.28.1
+    # via
+    #   anthropic
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   langgraph-sdk
+    #   langsmith
+    #   openai
+    #   safehttpx
+httpx-sse==0.4.3
+    # via langchain-community
+huggingface-hub==1.1.7
+    # via
+    #   gradio
+    #   gradio-client
+idna==3.11
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+    #   yarl
+jinja2==3.1.6
+    # via gradio
+jiter==0.12.0
+    # via
+    #   anthropic
+    #   openai
+jsonpatch==1.33
+    # via langchain-core
+jsonpointer==3.0.0
+    # via jsonpatch
+langchain==1.1.0
+    # via litigation-research (pyproject.toml)
+langchain-anthropic==1.2.0
+    # via litigation-research (pyproject.toml)
+langchain-classic==1.0.0
+    # via langchain-community
+langchain-community==0.4.1
+    # via litigation-research (pyproject.toml)
+langchain-core==1.1.0
+    # via
+    #   langchain
+    #   langchain-anthropic
+    #   langchain-classic
+    #   langchain-community
+    #   langchain-google-genai
+    #   langchain-openai
+    #   langchain-text-splitters
+    #   langgraph
+    #   langgraph-checkpoint
+    #   langgraph-prebuilt
+langchain-google-genai==3.2.0
+    # via litigation-research (pyproject.toml)
+langchain-openai==1.1.0
+    # via litigation-research (pyproject.toml)
+langchain-text-splitters==1.0.0
+    # via langchain-classic
+langgraph==1.0.4
+    # via langchain
+langgraph-checkpoint==3.0.1
+    # via
+    #   langgraph
+    #   langgraph-prebuilt
+langgraph-prebuilt==1.0.5
+    # via langgraph
+langgraph-sdk==0.2.12
+    # via langgraph
+langsmith==0.4.50
+    # via
+    #   langchain-classic
+    #   langchain-community
+    #   langchain-core
+markdown-it-py==4.0.0
+    # via rich
+markupsafe==3.0.3
+    # via
+    #   gradio
+    #   jinja2
+marshmallow==3.26.1
+    # via dataclasses-json
+mdurl==0.1.2
+    # via markdown-it-py
+multidict==6.7.0
+    # via
+    #   aiohttp
+    #   yarl
+mypy-extensions==1.1.0
+    # via typing-inspect
+numpy==2.3.5
+    # via
+    #   gradio
+    #   langchain-community
+    #   pandas
+openai==2.8.1
+    # via langchain-openai
+orjson==3.11.4
+    # via
+    #   gradio
+    #   langgraph-sdk
+    #   langsmith
+ormsgpack==1.12.0
+    # via langgraph-checkpoint
+packaging==25.0
+    # via
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   langchain-core
+    #   langsmith
+    #   marshmallow
+pandas==2.3.3
+    # via
+    #   litigation-research (pyproject.toml)
+    #   gradio
+pillow==12.0.0
+    # via gradio
+propcache==0.4.1
+    # via
+    #   aiohttp
+    #   yarl
+proto-plus==1.26.1
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+protobuf==6.33.1
+    # via
+    #   google-ai-generativelanguage
+    #   google-api-core
+    #   googleapis-common-protos
+    #   grpcio-status
+    #   proto-plus
+psycopg2-binary==2.9.11
+    # via litigation-research (pyproject.toml)
+pyasn1==0.6.1
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.4.2
+    # via google-auth
+pydantic==2.12.4
+    # via
+    #   anthropic
+    #   fastapi
+    #   gradio
+    #   langchain
+    #   langchain-anthropic
+    #   langchain-classic
+    #   langchain-core
+    #   langchain-google-genai
+    #   langgraph
+    #   langsmith
+    #   openai
+    #   pydantic-settings
+pydantic-core==2.41.5
+    # via pydantic
+pydantic-settings==2.12.0
+    # via langchain-community
+pydub==0.25.1
+    # via gradio
+pygments==2.19.2
+    # via rich
+pypdf==6.4.0
+    # via litigation-research (pyproject.toml)
+python-dateutil==2.9.0.post0
+    # via pandas
+python-dotenv==1.2.1
+    # via
+    #   litigation-research (pyproject.toml)
+    #   pydantic-settings
+python-multipart==0.0.20
+    # via gradio
+pytz==2025.2
+    # via pandas
+pyyaml==6.0.3
+    # via
+    #   gradio
+    #   huggingface-hub
+    #   langchain-classic
+    #   langchain-community
+    #   langchain-core
+regex==2025.11.3
+    # via tiktoken
+requests==2.32.5
+    # via
+    #   google-api-core
+    #   langchain-classic
+    #   langchain-community
+    #   langsmith
+    #   requests-toolbelt
+    #   tiktoken
+requests-toolbelt==1.0.0
+    # via langsmith
+rich==14.2.0
+    # via typer
+rsa==4.9.1
+    # via google-auth
+safehttpx==0.1.7
+    # via gradio
+semantic-version==2.10.0
+    # via gradio
+shellingham==1.5.4
+    # via
+    #   huggingface-hub
+    #   typer
+six==1.17.0
+    # via python-dateutil
+sniffio==1.3.1
+    # via
+    #   anthropic
+    #   openai
+sqlalchemy==2.0.44
+    # via
+    #   litigation-research (pyproject.toml)
+    #   langchain-classic
+    #   langchain-community
+starlette==0.50.0
+    # via
+    #   fastapi
+    #   gradio
+tenacity==9.1.2
+    # via
+    #   langchain-community
+    #   langchain-core
+tiktoken==0.12.0
+    # via langchain-openai
+tomlkit==0.13.3
+    # via gradio
+tqdm==4.67.1
+    # via
+    #   huggingface-hub
+    #   openai
+typer==0.20.0
+    # via gradio
+typer-slim==0.20.0
+    # via huggingface-hub
+typing-extensions==4.15.0
+    # via
+    #   anthropic
+    #   fastapi
+    #   gradio
+    #   gradio-client
+    #   grpcio
+    #   huggingface-hub
+    #   langchain-core
+    #   openai
+    #   pydantic
+    #   pydantic-core
+    #   sqlalchemy
+    #   typer
+    #   typer-slim
+    #   typing-inspect
+    #   typing-inspection
+typing-inspect==0.9.0
+    # via dataclasses-json
+typing-inspection==0.4.2
+    # via
+    #   pydantic
+    #   pydantic-settings
+tzdata==2025.2
+    # via pandas
+urllib3==2.5.0
+    # via requests
+uvicorn==0.38.0
+    # via gradio
+xxhash==3.6.0
+    # via langgraph
+yarl==1.22.0
+    # via aiohttp
+zstandard==0.25.0
+    # via langsmith

src/app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import pandas as pd
-from pipeline import LitigationPipeline
+from pipeline import LitigationPipeline, MODELS
+from db import init_db, DocumentLog
 
 
 # --- DEFAULT PROMPTS ---
@@ -43,6 +44,17 @@ def run_pipeline(files, api_key, model, context_limit, retries, sum_prompt, ver_
     
     return status_msg, logs_df
 
+def get_initial_logs():
+    """Fetch initial logs to populate the UI."""
+    db_session = init_db()
+    query = db_session.query(DocumentLog).order_by(DocumentLog.timestamp.desc()).limit(50)
+    df = pd.read_sql(query.statement, db_session.bind)
+    # Ensure columns match the expected headers in the UI
+    df_display = df.rename(columns={
+        "verification_score": "score",
+        "verifier_feedback": "feedback"
+    })
+    return df_display
 # --- GRADIO LAYOUT ---
 
 with gr.Blocks(title="Pharma Litigation AI Researcher") as demo:
@@ -62,7 +74,7 @@ with gr.Blocks(title="Pharma Litigation AI Researcher") as demo:
             gr.Markdown("---")
             gr.Markdown("### ⚙️ Configuration")
             api_key_input = gr.Textbox(label="API Key", type="password", placeholder="sk-...")
-            model_selector = gr.Dropdown(["google_genai:gemini-2.5-flash", "openai:gpt-4o", "anthropic:claude-sonnet-4-5-20250929"], label="Model", value="google_genai:gemini-2.5-flash")
+            model_selector = gr.Dropdown(list(MODELS.keys()), label="Model", value=list(MODELS.keys())[0])
             context_limit = gr.Number(label="Reset Context After (N docs)", value=5, precision=0)
             retry_limit = gr.Slider(minimum=0, maximum=5, value=2, step=1, label="Max Verification Retries")
             
@@ -73,7 +85,8 @@ with gr.Blocks(title="Pharma Litigation AI Researcher") as demo:
             with gr.Tabs():
                 with gr.TabItem("📊 Live Execution Logs"):
                     status_output = gr.Markdown("Waiting for input...")
-                    log_table = gr.Dataframe(label="Processing Logs (PostgreSQL)", headers=["filename", "status", "score", "feedback"])
+                    log_table = gr.Dataframe(label="Processing Logs (PostgreSQL)", 
+                                             headers=["filename", "status", "score", "feedback"])
 
     
 
@@ -91,6 +104,7 @@ with gr.Blocks(title="Pharma Litigation AI Researcher") as demo:
         ],
         outputs=[status_output, log_table]
     )
+    demo.load(get_initial_logs, None, log_table)
 
 if __name__ == "__main__":
     demo.launch()

src/pipeline.py

@@ -1,22 +1,34 @@
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
 from langchain.chat_models import init_chat_model
 from langchain_community.document_loaders import PyPDFLoader
+from pydantic import BaseModel
 import pandas as pd
 import os
 from db import init_db, DocumentLog
 
+MODELS = {
+    "Google Gemini 2.5 Flash":{"model_id":"google_genai:gemini-2.5-flash", "cost":0.00001},       
+    "OpenAI GPT-4o":{"model_id":"openai:gpt-4o", "cost":0.00003},
+    "Claude Sonnet 4.5":{"model_id":"anthropic:claude-sonnet-4-5-20250929", "cost":0.00002},
+}
+
+
+class VerifierResponse(BaseModel):
+    score: float
+    feedback: str   
+
+
 class LitigationPipeline:
-    def __init__(self, api_key, model_provider, batch_context_limit):
+    def __init__(self, api_key, model, batch_context_limit):
         self.api_key = api_key
-        self.provider = model_provider
+        self.model = model
         self.context_limit = int(batch_context_limit)
         self.processed_count = 0
         self.db_session = init_db()
         
         # Initialize LLM
         try:
-            self.llm = init_chat_model(model=model_provider, api_key=api_key)
+            self.llm = init_chat_model(model=MODELS[self.model]['model_id'], api_key=api_key)
         except Exception as e:
             raise ValueError("Unsupported Model: " + str(e))
 
@@ -26,7 +38,7 @@ class LitigationPipeline:
         # or clearing conversation buffers.
         print(f"--- SYSTEM: Context Reset Triggered after {self.processed_count} docs ---")
         # For this prototype, we re-initialize the LLM client to ensure no cached state remains
-        self.llm = init_chat_model(model=self.provider, api_key=self.api_key)
+        self.llm = init_chat_model(model=self.model, api_key=self.api_key)
 
     def extract_text(self, file_path):
         loader = PyPDFLoader(file_path)
@@ -43,8 +55,12 @@ class LitigationPipeline:
             ("system", system_instruction),
             ("user", f"Analyze this legal text:\n{text[:50000]}...") # Truncate for prototype safety
         ])
-        chain = prompt | self.llm | StrOutputParser()
-        return chain.invoke({})
+        chain = prompt | self.llm
+
+        result = chain.invoke({})
+        tokens = result.usage_metadata.get("total_tokens", 0)
+        cost =  tokens * MODELS[self.model]['cost'] 
+        return result.content, cost
 
     def verify(self, summary, original_text, verifier_prompt):
         """Critiques the summary. Returns (score, feedback)."""
@@ -52,19 +68,19 @@ class LitigationPipeline:
             ("system", verifier_prompt),
             ("user", f"Original Text (Excerpt): {original_text[:5000]}...\n\nSummary to Verify:\n{summary}")
         ])
-        chain = prompt | self.llm | StrOutputParser()
+        chain = prompt | self.llm.with_structured_output(VerifierResponse)
         result = chain.invoke({})
-        
+        cost =  len(result.feedback.split(" ")) * 0.7 * MODELS[self.model]['cost'] 
+
         # Simple parsing for prototype: Expecting "SCORE: X/10\nFEEDBACK: ..."
         try:
-            score_line = [line for line in result.split('\n') if "SCORE:" in line][0]
-            score = float(score_line.split(":")[1].strip().split("/")[0])
-            feedback = result.split("FEEDBACK:")[1].strip()
+            score = result.score
+            feedback = result.feedback
         except:
             score = 0
             feedback = "Parsing Error in Verifier Output"
             
-        return score, feedback
+        return score, feedback, cost
 
     def process_batch(self, file_paths, summary_prompt, verifier_prompt, retry_limit):
         results = []
@@ -84,7 +100,7 @@ class LitigationPipeline:
             final_score = 0
             
             extracted_text = ""
-
+            total_cost=0
             try:
                 extracted_text = self.extract_text(file_path)
                 
@@ -93,11 +109,11 @@ class LitigationPipeline:
                     print(f"Processing {file_name}: Attempt {attempts + 1}")
                     
                     # 1. Summarize
-                    current_summary = self.summarize(extracted_text, summary_prompt, feedback=current_feedback if attempts > 0 else None)
-                    
+                    current_summary, summary_cost = self.summarize(extracted_text, summary_prompt, feedback=current_feedback if attempts > 0 else None)
+                    total_cost += summary_cost 
                     # 2. Verify
-                    final_score, current_feedback = self.verify(current_summary, extracted_text, verifier_prompt)
-                    
+                    final_score, current_feedback, verifier_cost = self.verify(current_summary, extracted_text, verifier_prompt)
+                    total_cost += verifier_cost
                     # 3. Decision
                     if final_score >= 8.0: # Threshold for success
                         success = True
@@ -118,8 +134,8 @@ class LitigationPipeline:
                 verification_score=final_score,
                 verifier_feedback=current_feedback,
                 final_summary=current_summary,
-                model_used=self.provider,
-                cost_estimate=0.01 * (attempts + 1) # Placeholder logic
+                model_used=self.model,
+                cost_estimate=total_cost
             )
             self.db_session.add(log_entry)
             self.db_session.commit()