Create llm_interface.py

llm_interface.py

@@ -0,0 +1,531 @@
+"""
+Interface for interacting with Anthropic Claude API for:
+1. Extracting structured data from document sections
+2. Generating content for authoring
+3. Answering questions about documents via RAG
+"""
+
+import os
+import json
+import anthropic
+from typing import Dict, List, Any, Optional, Union
+import time
+
+class LLMInterface:
+    """Interface for interacting with LLMs, specifically Claude."""
+    
+    def __init__(self, api_key=None):
+        """Initialize the interface with an API key."""
+        if api_key:
+            self.api_key = api_key
+        else:
+            # Get from environment variable
+            self.api_key = os.environ.get("ANTHROPIC_API_KEY")
+            
+        if not self.api_key:
+            raise ValueError("Anthropic API Key is required")
+            
+        self.client = anthropic.Anthropic(api_key=self.api_key)
+    
+    def _call_claude(self, prompt: str, system: str = None, max_tokens: int = 4000, 
+                    temperature: float = 0.2, model: str = "claude-3-sonnet-20240229") -> str:
+        """
+        Make a call to Claude API.
+        
+        Args:
+            prompt: The prompt to send to Claude
+            system: Optional system prompt
+            max_tokens: Maximum tokens in the response
+            temperature: Temperature setting (0-1)
+            model: Model to use
+            
+        Returns:
+            Claude's response as a string
+        """
+        try:
+            messages = [{"role": "user", "content": prompt}]
+            
+            response = self.client.messages.create(
+                model=model,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                system=system,
+                messages=messages
+            )
+            
+            return response.content[0].text
+        except Exception as e:
+            print(f"Error calling Claude API: {e}")
+            # Wait and retry once on rate limiting
+            if "rate" in str(e).lower() or "timeout" in str(e).lower():
+                print("Rate limit hit, waiting 5 seconds...")
+                time.sleep(5)
+                try:
+                    response = self.client.messages.create(
+                        model=model,
+                        max_tokens=max_tokens,
+                        temperature=temperature,
+                        system=system,
+                        messages=messages
+                    )
+                    return response.content[0].text
+                except Exception as retry_e:
+                    print(f"Retry failed: {retry_e}")
+                    return f"Error: {retry_e}"
+            return f"Error: {e}"
+    
+    def _parse_json_from_response(self, response: str) -> Dict:
+        """
+        Extract and parse JSON from Claude's response.
+        
+        Args:
+            response: Claude's text response
+            
+        Returns:
+            Parsed JSON as a dictionary
+        """
+        try:
+            # Find JSON in the response (it might be wrapped in ```json or just be part of the text)
+            json_start = response.find('{')
+            json_end = response.rfind('}') + 1
+            
+            if json_start >= 0 and json_end > json_start:
+                json_str = response[json_start:json_end]
+                return json.loads(json_str)
+            else:
+                print("No JSON found in response")
+                return {}
+        except json.JSONDecodeError as e:
+            print(f"Error parsing JSON: {e}")
+            print(f"Response was: {response}")
+            return {}
+    
+    def extract_study_info(self, protocol_text: str) -> Dict:
+        """
+        Extract basic study information from protocol text.
+        
+        Args:
+            protocol_text: Text from the protocol
+            
+        Returns:
+            Dictionary with study information
+        """
+        system = """
+        You are an expert in clinical trial protocols with the specific task of extracting 
+        structured data from protocol text. Extract only the information that is explicitly 
+        stated in the text. If information is not available, use null or empty strings. 
+        Return a valid JSON object.
+        """
+        
+        prompt = """
+        Extract the following study information from the provided protocol text.
+        Return a valid JSON object with these keys:
+        {
+            "protocol_id": "string",       // The protocol identifier/number
+            "title": "string",             // The full protocol title
+            "phase": "string",             // Clinical trial phase
+            "status": "string",            // Protocol status if mentioned
+            "design_type": "string",       // Study design description (e.g., "Randomized, Double-Blind...")
+            "compound_id": "string",       // Investigational product identifier/name
+            "indication": "string",        // Disease or condition being studied
+            "planned_enrollment": "string" // Number of planned subjects/participants
+        }
+        
+        Protocol text:
+        """
+        
+        response = self._call_claude(prompt + protocol_text[:20000], system=system)
+        return self._parse_json_from_response(response)
+    
+    def extract_objectives_and_endpoints(self, section_text: str, protocol_id: str) -> Dict:
+        """
+        Extract objectives and their corresponding endpoints from protocol text.
+        
+        Args:
+            section_text: Text from the objectives/endpoints section
+            protocol_id: Protocol ID for reference
+            
+        Returns:
+            Dictionary with objectives and endpoints
+        """
+        system = """
+        You are an expert in clinical trial protocols with the specific task of extracting 
+        structured data about objectives and endpoints. Extract only the information that 
+        is explicitly stated in the text. Return the data as a valid JSON object.
+        """
+        
+        prompt = f"""
+        Extract the objectives and endpoints from the following protocol section text.
+        The protocol ID is: {protocol_id}
+        
+        Return a valid JSON object with these keys:
+        {{
+            "objectives": [
+                {{
+                    "type": "string",        // "Primary", "Secondary", or "Exploratory"
+                    "description": "string", // The full text description of the objective
+                    "id": "string"           // A generated identifier (e.g., "OBJ1", "OBJ2")
+                }}
+            ],
+            "endpoints": [
+                {{
+                    "type": "string",        // "Primary", "Secondary", or "Exploratory" 
+                    "name": "string",        // Short name of the endpoint
+                    "definition": "string",  // Full definition
+                    "objective_id": "string" // Reference to which objective this endpoint measures (if clear)
+                }}
+            ]
+        }}
+        
+        Section text:
+        """
+        
+        response = self._call_claude(prompt + section_text, system=system)
+        return self._parse_json_from_response(response)
+    
+    def extract_population_criteria(self, section_text: str, protocol_id: str) -> Dict:
+        """
+        Extract inclusion and exclusion criteria from protocol text.
+        
+        Args:
+            section_text: Text from the population/criteria section
+            protocol_id: Protocol ID for reference
+            
+        Returns:
+            Dictionary with inclusion and exclusion criteria
+        """
+        system = """
+        You are an expert in clinical trial protocols with the specific task of extracting 
+        structured data about inclusion and exclusion criteria. Extract the criteria 
+        exactly as stated in the text, preserving numbering and formatting. Return the 
+        data as a valid JSON object.
+        """
+        
+        prompt = f"""
+        Extract the inclusion and exclusion criteria from the following protocol section.
+        The protocol ID is: {protocol_id}
+        
+        Return a valid JSON object with these keys:
+        {{
+            "inclusion_criteria": [
+                {{
+                    "number": number or null,  // The criterion number if available (e.g., 1, 2)
+                    "text": "string",          // The full text of the criterion
+                    "attribute": "string",     // The characteristic being evaluated, if clear (e.g., "Age", "BMI")
+                    "operator": "string",      // The comparison operator if applicable (e.g., ">", "<", "=")
+                    "value": "string"          // The threshold value if applicable (e.g., "18 years")
+                }}
+            ],
+            "exclusion_criteria": [
+                {{
+                    "number": number or null,
+                    "text": "string",
+                    "attribute": "string",
+                    "operator": "string",
+                    "value": "string"
+                }}
+            ]
+        }}
+        
+        Section text:
+        """
+        
+        response = self._call_claude(prompt + section_text, system=system)
+        return self._parse_json_from_response(response)
+    
+    def extract_study_design(self, section_text: str, protocol_id: str) -> Dict:
+        """
+        Extract study design information from protocol text.
+        
+        Args:
+            section_text: Text from the study design section
+            protocol_id: Protocol ID for reference
+            
+        Returns:
+            Dictionary with study design information
+        """
+        system = """
+        You are an expert in clinical trial protocols with the specific task of extracting 
+        structured data about study design. Extract only information that is explicitly 
+        stated in the text. Return the data as a valid JSON object.
+        """
+        
+        prompt = f"""
+        Extract the study design information from the following protocol section.
+        The protocol ID is: {protocol_id}
+        
+        Return a valid JSON object with these keys:
+        {{
+            "design_type": "string",          // E.g., "Randomized, Double-blind, Placebo-controlled"
+            "study_parts": [                  // List of different parts/cohorts if applicable
+                {{
+                    "part": "string",           // Identifier (e.g., "Part A", "Cohort 1")
+                    "description": "string",    // Description
+                    "population": "string",     // E.g., "Healthy Volunteers" or "T2DM Patients"
+                    "planned_n": "string"       // Planned number of subjects
+                }}
+            ],
+            "randomization": "string",        // Description of randomization process
+            "blinding": "string",             // Description of blinding (e.g., "Double-blind")
+            "duration": "string",             // Study duration information
+            "dose_info": "string"             // Information about dosing if mentioned
+        }}
+        
+        Section text:
+        """
+        
+        response = self._call_claude(prompt + section_text, system=system)
+        return self._parse_json_from_response(response)
+    
+    def extract_statistical_methods(self, section_text: str, protocol_id: str) -> Dict:
+        """
+        Extract statistical analysis methods from SAP or protocol text.
+        
+        Args:
+            section_text: Text from the statistical methods section
+            protocol_id: Protocol ID for reference
+            
+        Returns:
+            Dictionary with statistical methods information
+        """
+        system = """
+        You are an expert in clinical trial statistics with the specific task of extracting 
+        structured data about statistical methods from protocols or SAPs. Return the data 
+        as a valid JSON object.
+        """
+        
+        prompt = f"""
+        Extract the statistical methods information from the following section.
+        The protocol ID is: {protocol_id}
+        
+        Return a valid JSON object with these keys:
+        {{
+            "analysis_populations": [
+                {{
+                    "name": "string",        // E.g., "Full Analysis Set", "Safety Population"
+                    "definition": "string"   // Definition of the population
+                }}
+            ],
+            "primary_analysis": {{
+                "endpoint": "string",        // Primary endpoint being analyzed
+                "method": "string",          // Statistical method (e.g., "MMRM", "t-test")
+                "covariates": ["string"],    // List of covariates if mentioned
+                "handling_missing": "string" // How missing data is handled
+            }},
+            "secondary_analyses": [
+                {{
+                    "endpoint": "string",
+                    "method": "string",
+                    "covariates": ["string"],
+                    "handling_missing": "string"
+                }}
+            ],
+            "multiplicity": "string",        // How multiplicity is addressed
+            "sample_size_justification": "string" // Sample size rationale
+        }}
+        
+        Section text:
+        """
+        
+        response = self._call_claude(prompt + section_text, system=system)
+        return self._parse_json_from_response(response)
+    
+    def extract_assessments(self, section_text: str, protocol_id: str) -> Dict:
+        """
+        Extract assessment information from protocol text.
+        
+        Args:
+            section_text: Text from the assessments section
+            protocol_id: Protocol ID for reference
+            
+        Returns:
+            Dictionary with assessment information
+        """
+        system = """
+        You are an expert in clinical trial protocols with the specific task of extracting 
+        structured data about assessments and procedures. Return the data as a valid JSON object.
+        """
+        
+        prompt = f"""
+        Extract information about assessments and procedures from the following protocol section.
+        The protocol ID is: {protocol_id}
+        
+        Return a valid JSON object with these keys:
+        {{
+            "assessments": [
+                {{
+                    "name": "string",         // Name of assessment (e.g., "OGTT", "ECG")
+                    "type": "string",         // Type (e.g., "Safety", "PK", "PD")
+                    "description": "string",  // Description of the procedure
+                    "timing": "string",       // When it's performed
+                    "analytes": ["string"]    // Measured analytes if applicable
+                }}
+            ]
+        }}
+        
+        Section text:
+        """
+        
+        response = self._call_claude(prompt + section_text, system=system)
+        return self._parse_json_from_response(response)
+    
+    def generate_content_from_knowledge(self, section_type: str, context: List[Dict], 
+                                       protocol_id: str = None, style_guide: str = None) -> str:
+        """
+        Generate document content based on knowledge extracted from similar documents.
+        
+        Args:
+            section_type: Type of section to generate (e.g., "Introduction", "Study Design")
+            context: List of relevant text chunks from knowledge base
+            protocol_id: Optional protocol ID for reference
+            style_guide: Optional style guide instructions
+            
+        Returns:
+            Generated content as a string
+        """
+        system = """
+        You are an expert medical writer who specializes in pharmaceutical R&D documents
+        like protocols, SAPs, and CSRs. Your task is to draft high-quality content
+        based on similar examples, following the conventions of scientific/medical writing
+        and any provided style guides.
+        """
+        
+        # Prepare context text
+        context_text = ""
+        for i, chunk in enumerate(context):
+            context_text += f"\nEXAMPLE {i+1} (Source: {chunk.get('metadata', {}).get('source', 'Unknown')})\n"
+            context_text += chunk.get('page_content', '')
+            context_text += "\n" + "-"*50 + "\n"
+        
+        protocol_ref = f"for protocol {protocol_id}" if protocol_id else ""
+        style_instructions = f"\nFollow these style guidelines:\n{style_guide}" if style_guide else ""
+        
+        prompt = f"""
+        Please draft a {section_type} section {protocol_ref} for a clinical study document.
+        
+        The content should be:
+        1. Well-structured and professionally written
+        2. Scientifically accurate and precise
+        3. Appropriate for a regulatory/scientific audience
+        4. In line with typical conventions for pharmaceutical documents{style_instructions}
+        
+        Here are examples of similar content from other documents to guide your writing:
+        {context_text}
+        
+        Please draft a complete {section_type} section that follows these examples in style and 
+        structure but is original.
+        """
+        
+        # Use a higher max tokens for content generation
+        response = self._call_claude(prompt, system=system, max_tokens=4000, temperature=0.3)
+        return response
+    
+    def answer_protocol_question(self, question: str, context: List[Dict], 
+                               chat_history: List[Dict] = None) -> str:
+        """
+        Answer a question about protocols using retrieved context.
+        
+        Args:
+            question: User's question
+            context: List of relevant text chunks from knowledge base
+            chat_history: Optional list of previous interactions
+            
+        Returns:
+            Answer as a string
+        """
+        system = """
+        You are a Protocol Coach, an expert assistant specializing in pharmaceutical R&D documents.
+        Your role is to answer questions about clinical study protocols, SAPs, and other related documents
+        using the specific context provided. Base your answers strictly on the provided context and
+        indicate when information might not be available in the provided excerpts.
+        
+        Always cite the source documents when answering questions.
+        """
+        
+        # Prepare context text
+        context_text = ""
+        for i, chunk in enumerate(context):
+            source = chunk.get('metadata', {}).get('source', 'Unknown')
+            section = chunk.get('metadata', {}).get('section', 'Unknown section')
+            context_text += f"\nCONTEXT {i+1} [Source: {source}, Section: {section}]\n"
+            context_text += chunk.get('page_content', '')
+            context_text += "\n" + "-"*50 + "\n"
+        
+        # Prepare chat history if available
+        history_text = ""
+        if chat_history and len(chat_history) > 0:
+            history_text = "\nPrevious conversation:\n"
+            for entry in chat_history[-3:]:  # Only use last 3 exchanges for context
+                if 'user' in entry:
+                    history_text += f"User: {entry['user']}\n"
+                if 'assistant' in entry:
+                    history_text += f"Assistant: {entry['assistant']}\n"
+            history_text += "\n"
+        
+        prompt = f"""
+        {history_text}
+        User question: {question}
+        
+        Please answer the question based on the following context from clinical documents:
+        {context_text}
+        
+        Answer the question comprehensively using only the information in the provided context.
+        If the context doesn't contain sufficient information to provide a complete answer,
+        clearly state which aspects you can and cannot address based on the available information.
+        """
+        
+        response = self._call_claude(prompt, system=system, max_tokens=2000, temperature=0.2)
+        return response
+    
+    def find_document_connections(self, source_doc_info: Dict, target_doc_info: Dict,
+                                entity_pairs: List[Dict]) -> str:
+        """
+        Analyze connections between two documents based on entity pairs.
+        
+        Args:
+            source_doc_info: Information about the source document
+            target_doc_info: Information about the target document
+            entity_pairs: List of potentially matching entities from both documents
+            
+        Returns:
+            Analysis of connections as a string
+        """
+        system = """
+        You are an expert in pharmaceutical R&D document analysis, specialized in 
+        identifying relationships, consistency, and traceability between related 
+        documents like protocols and SAPs. Your task is to analyze potential 
+        matches between entities in different documents and assess their alignment.
+        """
+        
+        # Convert entity pairs to formatted text
+        entity_pairs_text = ""
+        for i, pair in enumerate(entity_pairs):
+            entity_pairs_text += f"\nCOMPARISON {i+1}:\n"
+            entity_pairs_text += f"Source: {pair.get('source_text', 'Not available')}\n"
+            entity_pairs_text += f"Target: {pair.get('target_text', 'Not available')}\n"
+            entity_pairs_text += f"Entity Type: {pair.get('entity_type', 'Unknown')}\n"
+            entity_pairs_text += "-"*50 + "\n"
+        
+        prompt = f"""
+        Analyze the connections between these two pharmaceutical documents:
+        
+        SOURCE DOCUMENT: {source_doc_info.get('title', 'Unknown')} (Type: {source_doc_info.get('type', 'Unknown')})
+        TARGET DOCUMENT: {target_doc_info.get('title', 'Unknown')} (Type: {target_doc_info.get('type', 'Unknown')})
+        
+        I'll provide pairs of potentially related elements from both documents. For each pair, assess:
+        1. Whether they refer to the same entity or concept
+        2. The level of consistency between them (High/Medium/Low)
+        3. Any notable differences or potential issues
+        
+        Here are the element pairs to analyze:
+        {entity_pairs_text}
+        
+        Provide:
+        1. A summary of the overall consistency between documents
+        2. Specific observations about each compared element
+        3. Potential implications of any inconsistencies
+        4. Recommendations for improving alignment
+        """
+        
+        response = self._call_claude(prompt, system=system, max_tokens=3000, temperature=0.2)
+        return response