Spaces:
Sleeping
Sleeping
| import os | |
| import requests | |
| import json | |
| import logging | |
| import time | |
| from typing import List, Dict, Tuple | |
| logger = logging.getLogger(__name__) | |
| class NVIDIALLamaClient: | |
| def __init__(self): | |
| self.api_key = os.getenv("NVIDIA_URI") | |
| if not self.api_key: | |
| raise ValueError("NVIDIA_URI environment variable not set") | |
| # Correct NVIDIA Integrate API base | |
| self.base_url = "https://integrate.api.nvidia.com/v1" | |
| self.model = "meta/llama-3.1-8b-instruct" | |
| def generate_keywords(self, user_query: str) -> List[str]: | |
| """Use Llama to generate search keywords from user query""" | |
| try: | |
| prompt = f"""Given this medical question: "{user_query}" | |
| Generate 3-5 specific search keywords that would help find relevant medical information online. | |
| Focus on medical terms, symptoms, conditions, treatments, or procedures mentioned. | |
| Return only the keywords separated by commas, no explanations. | |
| Keywords:""" | |
| response = self._call_llama(prompt) | |
| # Extract keywords from response | |
| keywords = [kw.strip() for kw in response.split(',') if kw.strip()] | |
| logger.info(f"Generated keywords: {keywords}") | |
| return keywords[:5] # Limit to 5 keywords | |
| except Exception as e: | |
| logger.error(f"Failed to generate keywords: {e}") | |
| return [user_query] # Fallback to original query | |
| def summarize_documents(self, documents: List[Dict], user_query: str) -> Tuple[str, Dict[int, str]]: | |
| """Use Llama to summarize documents and return summary with URL mapping""" | |
| try: | |
| # Import summarizer here to avoid circular imports | |
| from summarizer import summarizer | |
| # Use the summarizer for document summarization | |
| combined_summary, url_mapping = summarizer.summarize_documents(documents, user_query) | |
| return combined_summary, url_mapping | |
| except Exception as e: | |
| logger.error(f"Failed to summarize documents: {e}") | |
| return "", {} | |
| def _call_llama(self, prompt: str, max_retries: int = 3) -> str: | |
| """Make API call to NVIDIA Llama model with retry logic""" | |
| for attempt in range(max_retries): | |
| try: | |
| headers = { | |
| "Authorization": f"Bearer {self.api_key}", | |
| "Content-Type": "application/json" | |
| } | |
| payload = { | |
| "model": self.model, | |
| "messages": [ | |
| { | |
| "role": "user", | |
| "content": prompt | |
| } | |
| ], | |
| "temperature": 0.7, | |
| "max_tokens": 1000 | |
| } | |
| response = requests.post( | |
| f"{self.base_url}/chat/completions", | |
| headers=headers, | |
| json=payload, | |
| timeout=30 | |
| ) | |
| response.raise_for_status() | |
| result = response.json() | |
| content = result['choices'][0]['message']['content'].strip() | |
| if not content: | |
| raise ValueError("Empty response from Llama API") | |
| return content | |
| except requests.exceptions.Timeout: | |
| logger.warning(f"Llama API timeout (attempt {attempt + 1}/{max_retries})") | |
| if attempt == max_retries - 1: | |
| raise | |
| time.sleep(2 ** attempt) # Exponential backoff | |
| except requests.exceptions.RequestException as e: | |
| logger.warning(f"Llama API request failed (attempt {attempt + 1}/{max_retries}): {e}") | |
| if attempt == max_retries - 1: | |
| raise | |
| time.sleep(2 ** attempt) | |
| except Exception as e: | |
| logger.error(f"Llama API call failed: {e}") | |
| raise | |
| def process_search_query(user_query: str, search_results: List[Dict]) -> Tuple[str, Dict[int, str]]: | |
| """Process search results using Llama model""" | |
| try: | |
| llama_client = NVIDIALLamaClient() | |
| # Generate search keywords | |
| keywords = llama_client.generate_keywords(user_query) | |
| # Summarize documents | |
| summary, url_mapping = llama_client.summarize_documents(search_results, user_query) | |
| return summary, url_mapping | |
| except Exception as e: | |
| logger.error(f"Failed to process search query: {e}") | |
| return "", {} |