Update app.py

app.py

@@ -1,23 +1,24 @@
-import streamlit as st
 import os
+import time
+import json
+import numpy as np
+import streamlit as st
 import tempfile
 import requests
-import json
+from dotenv import load_dotenv
 from openai import OpenAI
 from rdkit import Chem
-from rdkit.Chem import Draw
+from rdkit.Chem import Draw, Descriptors, Crippen, Lipinski
 import faiss
-import numpy as np
 from PyPDF2 import PdfReader
-from typing import List, Dict, Optional, Tuple
-from dotenv import load_dotenv
-from fpdf import FPDF
+from typing import List, Dict, Tuple
 
-# --------------------------
-# Configuration & Environment
-# --------------------------
+# Load environment variables
 load_dotenv()
 
+# --------------------------
+# Configuration Settings
+# --------------------------
 class AppConfig:
     OPENAI_MODEL = "gpt-4-turbo-preview"
     EMBEDDING_MODEL = "text-embedding-3-large"
@@ -26,13 +27,11 @@ class AppConfig:
     RAG_THRESHOLD = 0.78
     MAX_CONTEXT_CHUNKS = 5
 
-# --------------------------
-# OpenAI Client Initialization
-# --------------------------
+# Initialize OpenAI Client
 client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
 # --------------------------
-# RAG Vector Index Implementation
+# Vector Index for RAG
 # --------------------------
 class VectorIndex:
     def __init__(self):
@@ -41,109 +40,64 @@ class VectorIndex:
         self.metadata = []
 
     def add_document(self, text: str, source: str):
-        """Add document to vector index with metadata"""
         embedding = self._get_embedding(text)
         self.index.add(np.array([embedding], dtype=np.float32))
         self.documents.append(text)
-        self.metadata.append({
-            "source": source,
-            "timestamp": time.time()
-        })
+        self.metadata.append({"source": source, "timestamp": time.time()})
 
     def search(self, query: str) -> List[Tuple[str, Dict]]:
-        """Search index with relevance scoring"""
         query_embed = self._get_embedding(query)
         distances, indices = self.index.search(np.array([query_embed], dtype=np.float32), AppConfig.MAX_CONTEXT_CHUNKS)
         
-        results = []
-        for i, idx in enumerate(indices[0]):
-            if idx >= 0 and distances[0][i] < AppConfig.RAG_THRESHOLD:
-                results.append((
-                    self.documents[idx],
-                    {**self.metadata[idx], "score": float(distances[0][i])}
-                ))
-        return results
+        return [
+            (self.documents[idx], {**self.metadata[idx], "score": float(distances[0][i])})
+            for i, idx in enumerate(indices[0]) if idx >= 0 and distances[0][i] < AppConfig.RAG_THRESHOLD
+        ]
 
     def _get_embedding(self, text: str) -> List[float]:
-        """Generate embeddings using OpenAI"""
-        response = client.embeddings.create(
-            input=text,
-            model=AppConfig.EMBEDDING_MODEL
-        )
+        response = client.embeddings.create(input=text, model=AppConfig.EMBEDDING_MODEL)
         return response.data[0].embedding
 
-# Initialize vector index
 knowledge_base = VectorIndex()
 
 # --------------------------
-# Pharmaceutical Tools with OpenAI Function Calling
+# Pharmaceutical Tools
 # --------------------------
 class PharmaTools:
     @staticmethod
     def clinical_trial_search(query: str) -> Dict:
-        """Search clinical trials with safety checks"""
         try:
-            response = client.chat.completions.create(
-                model=AppConfig.OPENAI_MODEL,
-                messages=[{
-                    "role": "user",
-                    "content": f"Generate clinicaltrials.gov API parameters for: {query}"
-                }],
-                tools=[{
-                    "type": "function",
-                    "function": {
-                        "name": "clinical_trial_search",
-                        "description": "Search clinical trials database",
-                        "parameters": {
-                            "type": "object",
-                            "properties": {
-                                "query": {"type": "string"},
-                                "max_results": {"type": "integer"}
-                            }
-                        }
-                    }
-                }]
-            )
-            
-            # Execute actual API call here
-            return {"results": "Sample clinical trial data"}
-        
+            return {"results": "Sample clinical trial data"}  # Placeholder for API integration
         except Exception as e:
             return {"error": str(e)}
 
     @staticmethod
     def molecular_analysis(smiles: str) -> Dict:
-        """Analyze molecular structure with RDKit"""
         try:
             mol = Chem.MolFromSmiles(smiles)
             if not mol:
                 return {"error": "Invalid SMILES"}
             
             properties = {
-                "molecular_weight": Chem.Descriptors.ExactMolWt(mol),
-                "logp": Chem.Crippen.MolLogP(mol),
-                "h_bond_donors": Chem.Lipinski.NumHDonors(mol)
+                "molecular_weight": Descriptors.ExactMolWt(mol),
+                "logp": Crippen.MolLogP(mol),
+                "h_bond_donors": Lipinski.NumHDonors(mol)
             }
             
-            # Add AI-powered predictions
             ai_analysis = client.chat.completions.create(
                 model=AppConfig.OPENAI_MODEL,
-                messages=[{
-                    "role": "user",
-                    "content": f"Predict pharmaceutical properties for SMILES {smiles}:\n{properties}"
-                }]
+                messages=[{"role": "user", "content": f"Predict properties for SMILES {smiles}:\n{properties}"}]
             )
             
             return {
                 "calculated": properties,
                 "ai_predictions": json.loads(ai_analysis.choices[0].message.content)
             }
-            
         except Exception as e:
             return {"error": str(e)}
 
 # --------------------------
-# AI Agent with RAG Integration
+# Pharma AI Agent
 # --------------------------
 class PharmaAgent:
     def __init__(self):
@@ -151,112 +105,50 @@ class PharmaAgent:
         self.system_prompt = self._build_system_prompt()
 
     def query(self, prompt: str) -> str:
-        """Execute query with RAG context"""
-        # Retrieve relevant knowledge
         rag_context = knowledge_base.search(prompt)
-        
-        # Build context-aware prompt
-        messages = [{
-            "role": "system",
-            "content": self.system_prompt
-        }]
+        messages = [{"role": "system", "content": self.system_prompt}]
         
         if rag_context:
             messages.append({
                 "role": "assistant",
-                "content": "Relevant knowledge:\n" + "\n".join([f"[Source: {meta['source']}]\n{text}" 
-                              for text, meta in rag_context])
+                "content": "Relevant knowledge:\n" + "\n".join([f"[Source: {meta['source']}]\n{text}" for text, meta in rag_context])
             })
-            
-        messages.append({
-            "role": "user",
-            "content": prompt
-        })
         
-        # Execute OpenAI query
+        messages.append({"role": "user", "content": prompt})
+        
         try:
             response = client.chat.completions.create(
                 model=AppConfig.OPENAI_MODEL,
                 messages=messages,
-                tools=self._tools_schema(),
                 timeout=AppConfig.API_TIMEOUT
             )
-            
-            return self._process_response(response)
-            
+            return response.choices[0].message.content
         except Exception as e:
             return f"Error: {str(e)}"
 
     def _build_system_prompt(self) -> str:
-        """Construct dynamic system prompt"""
-        return f"""You are a pharmaceutical research AI with access to:
-        - Molecular analysis tools
-        - Clinical trial databases
-        - Latest research via RAG (Updated: {time.ctime()})
-        
-        Follow these rules:
-        1. Always validate chemical structures
-        2. Cite sources from RAG context
-        3. Check for recent regulatory updates
-        4. Maintain safety protocols"""
-
-    def _tools_schema(self) -> List[Dict]:
-        """Generate OpenAI-compatible tool schema"""
-        return [{
-            "type": "function",
-            "function": {
-                "name": tool.__name__,
-                "description": tool.__doc__,
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        # Add parameter definitions here
-                    }
-                }
-            }
-        } for tool in self.tools]
-
-    def _process_response(self, response) -> str:
-        """Handle tool calls and response parsing"""
-        # Add tool execution logic here
-        return response.choices[0].message.content
+        return f"""You are a pharmaceutical AI with access to molecular analysis, clinical trial data, and research via RAG (Updated: {time.ctime()}). Follow safety protocols."""
 
 # --------------------------
-# Streamlit UI Components
+# Streamlit UI
 # --------------------------
 def main():
-    st.set_page_config(
-        page_title="PharmaAI Research Suite",
-        page_icon="🧬",
-        layout="wide"
-    )
-    
+    st.set_page_config(page_title="PharmaAI Research Suite", page_icon="🧬", layout="wide")
     st.title("PharmaAI Research Suite")
     st.markdown("Integrated AI Platform for Pharmaceutical Research")
     
-    # Initialize components
     agent = PharmaAgent()
     
-    # Knowledge Base Management
     with st.sidebar:
         st.header("🧠 Knowledge Base")
-        uploaded_files = st.file_uploader("Upload Research Documents", 
-                                        type=["pdf", "txt"],
-                                        accept_multiple_files=True)
+        uploaded_files = st.file_uploader("Upload Research Documents", type=["pdf", "txt"], accept_multiple_files=True)
         
         if uploaded_files:
             for file in uploaded_files:
-                text = ""
-                if file.type == "application/pdf":
-                    reader = PdfReader(file)
-                    text = "\n".join([page.extract_text() for page in reader.pages])
-                else:
-                    text = file.getvalue().decode()
-                
+                text = "\n".join([page.extract_text() for page in PdfReader(file).pages]) if file.type == "application/pdf" else file.getvalue().decode()
                 knowledge_base.add_document(text, file.name)
             st.success(f"Added {len(uploaded_files)} documents to knowledge base")
     
-    # Main Interface
     tab1, tab2, tab3 = st.tabs(["Drug Development", "Molecular Analysis", "Literature Review"])
     
     with tab1:
@@ -265,13 +157,7 @@ def main():
         strategy = st.selectbox("Development Strategy", ["First-in-class", "Me-too", "Biologic"])
         
         if st.button("Generate Development Plan"):
-            prompt = f"""Develop {strategy} drug targeting {target} considering:
-            - Target validation
-            - Safety profile
-            - Competitive landscape
-            - Regulatory pathway"""
-            
-            response = agent.query(prompt)
+            response = agent.query(f"Develop {strategy} drug targeting {target}.")
             st.markdown(response)
     
     with tab2:
@@ -284,18 +170,11 @@ def main():
                 st.error(analysis["error"])
             else:
                 col1, col2 = st.columns(2)
-                with col1:
-                    st.subheader("Calculated Properties")
-                    st.json(analysis["calculated"])
-                with col2:
-                    st.subheader("AI Predictions")
-                    st.json(analysis["ai_predictions"])
-                
-                # Visualization
+                col1.json(analysis["calculated"])
+                col2.json(analysis["ai_predictions"])
                 mol = Chem.MolFromSmiles(smiles)
                 if mol:
-                    img = Draw.MolToImage(mol, size=(400, 400))
-                    st.image(img, caption="Molecular Structure")
+                    st.image(Draw.MolToImage(mol, size=(400, 400)), caption="Molecular Structure")
     
     with tab3:
         st.header("Literature Review")
@@ -303,12 +182,9 @@ def main():
         
         if research_query:
             results = knowledge_base.search(research_query)
-            if not results:
-                st.info("No relevant documents found")
-            else:
-                for text, meta in results:
-                    with st.expander(f"Source: {meta['source']} (Score: {meta['score']:.2f})"):
-                        st.markdown(f"```\n{text[:1000]}...\n```")
+            for text, meta in results:
+                with st.expander(f"Source: {meta['source']} (Score: {meta['score']:.2f})"):
+                    st.markdown(f"```\n{text[:1000]}...\n```")
 
 if __name__ == "__main__":
-    main()
+    main()