Spaces:
Running
Running
File size: 2,622 Bytes
409c17a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
"""
Infrastructure - Gemini LLM Service
"""
from typing import AsyncIterator, List
import google.generativeai as genai
from app.domain.interfaces import ILLM, LLMMessage, LLMResponse
class GeminiLLM(ILLM):
"""Gemini LLM implementation"""
def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
genai.configure(api_key=api_key)
self.model_name = model_name
self.model = genai.GenerativeModel(model_name)
async def generate(
self,
messages: List[LLMMessage],
temperature: float = 0.7,
max_tokens: int = 2048,
stream: bool = False,
) -> LLMResponse:
"""Generate response from Gemini"""
# Convert messages to Gemini format
prompt = self._build_prompt(messages)
# Generate
response = await self.model.generate_content_async(
prompt,
generation_config=genai.types.GenerationConfig(
temperature=temperature, max_output_tokens=max_tokens
),
)
# Count tokens (approximate)
tokens_used = len(prompt.split()) + len(response.text.split())
return LLMResponse(
content=response.text,
model=self.model_name,
tokens_used=tokens_used,
finish_reason="stop",
)
async def generate_stream(
self,
messages: List[LLMMessage],
temperature: float = 0.7,
max_tokens: int = 2048,
) -> AsyncIterator[str]:
"""Generate streaming response from Gemini"""
prompt = self._build_prompt(messages)
response = await self.model.generate_content_async(
prompt,
generation_config=genai.types.GenerationConfig(
temperature=temperature, max_output_tokens=max_tokens
),
stream=True,
)
async for chunk in response:
if chunk.text:
yield chunk.text
def get_model_name(self) -> str:
"""Get model name"""
return self.model_name
def _build_prompt(self, messages: List[LLMMessage]) -> str:
"""Build prompt from messages"""
parts = []
for msg in messages:
if msg.role == "system":
parts.append(f"System: {msg.content}")
elif msg.role == "user":
parts.append(f"User: {msg.content}")
elif msg.role == "assistant":
parts.append(f"Assistant: {msg.content}")
return "\n\n".join(parts)
|