llm.py

import os
import json
import google.generativeai as genai
from langchain.schema import AIMessage
from crisprgpt.safety import WARNING_PRIVACY, contains_identifiable_genes
from util import get_logger

logger = get_logger(__name__)

# Load API key from environment
GEMINI_API_KEY = "AIzaSyAnmXi2fL6hHwC4vjGYEA93QlR1VQN9Q9w"

genai.configure(api_key=GEMINI_API_KEY)


class IdentifiableGeneError(ValueError):
    """Raised when identifiable genes are detected in the input."""


class FakeChatGemini:
    """Debug class for simulating Gemini responses."""
    def __call__(self, inputs):
        logger.info("FakeChatGemini Called")
        response = input("Fake response: ")
        return AIMessage(content=response)


class GeminiChat:
    """Wrapper for Gemini API to behave like the old OpenAIChat class."""

    def __init__(self, model_name="gemini-1.5-flash"):
        self.model_name = model_name
        self.model = genai.GenerativeModel(model_name)

    def chat(self, request: str):
        """Handles structured JSON response requests."""
        if contains_identifiable_genes(request):
            raise IdentifiableGeneError(WARNING_PRIVACY)

        try:
            response = self.model.generate_content(request)
            if not response or not response.text:
                return {"error": "No response from Gemini."}

            raw_text = (
                response.text.strip()
                .removeprefix("```json")
                .removeprefix("```")
                .removesuffix("```")
                .strip()
            )

            try:
                return json.loads(raw_text)
            except json.JSONDecodeError:
                logger.warning("Gemini response was not valid JSON, returning raw text.")
                return {"response": raw_text}

        except Exception as e:
            logger.error(f"Gemini error: {e}")
            return {"error": str(e)}

    def QA(self, request: str):
        """Simple question-answering wrapper."""
        try:
            response = self.model.generate_content(request)
            return response.text if response and response.text else "No answer."
        except Exception as e:
            logger.error(f"QA error: {e}")
            return f"Error: {str(e)}"