44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
# [DEF:llm_core:Module]
|
|
import os
|
|
import logging
|
|
from typing import Type, TypeVar
|
|
from pydantic import BaseModel
|
|
|
|
# Pylance often fails to resolve dynamic exports in google-generativeai
|
|
import google.generativeai as genai # type: ignore
|
|
|
|
logger = logging.getLogger("LLM_Core")
|
|
T = TypeVar("T", bound=BaseModel)
|
|
|
|
class GeminiProcessor:
|
|
def __init__(self, model_name: str = "gemini-1.5-flash"):
|
|
api_key = os.getenv("GOOGLE_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("[FATAL] GOOGLE_API_KEY not found.")
|
|
|
|
# Explicit type ignore for Pylance strict mode
|
|
genai.configure(api_key=api_key) # type: ignore
|
|
self.model = genai.GenerativeModel(model_name) # type: ignore
|
|
|
|
def generate_structured(self, prompt: str, content: str, schema: Type[T]) -> T:
|
|
logger.info(f"[GeminiProcessor] Structured generation for {schema.__name__}")
|
|
full_prompt = f"{prompt}\n\nINPUT TEXT:\n{content}"
|
|
|
|
response = self.model.generate_content(
|
|
full_prompt,
|
|
generation_config=genai.GenerationConfig( # type: ignore
|
|
response_mime_type="application/json",
|
|
response_schema=schema
|
|
)
|
|
)
|
|
return schema.model_validate_json(response.text)
|
|
|
|
def generate_text(self, system_instruction: str, user_content: str) -> str:
|
|
# Re-instantiate model with system instruction
|
|
model_w_sys = genai.GenerativeModel( # type: ignore
|
|
self.model.model_name,
|
|
system_instruction=system_instruction
|
|
)
|
|
response = model_w_sys.generate_content(user_content)
|
|
return response.text
|
|
# [/DEF:llm_core] |