doc_knowledge_base / llm_interface.py
cryogenic22's picture
Create llm_interface.py
6ed286a verified
"""
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