Gemini Client

pip install graphiti-core[google-genai]

from graphiti_core.llm_client import GeminiClient
from graphiti_core.llm_client.config import LLMConfig
from graphiti_core.prompts.models import Message
from pydantic import BaseModel

# Initialize client
client = GeminiClient(
    config=LLMConfig(
        api_key="your-google-api-key",
        model="gemini-3-flash-preview",
        temperature=1.0
    )
)

# Define response structure
class Summary(BaseModel):
    title: str
    main_points: list[str]
    word_count: int

# Generate structured response
messages = [
    Message(role="system", content="Summarize the following article."),
    Message(role="user", content="Long article text...")
]

response = await client.generate_response(
    messages=messages,
    response_model=Summary
)

# Automatic: uses 65536 for gemini-3-flash-preview
client = GeminiClient(
    config=LLMConfig(model="gemini-3-flash-preview")
)

# Override: use 32K for all requests
client = GeminiClient(
    config=LLMConfig(model="gemini-2.5-pro"),
    max_tokens=32000
)

# Per-request: 16K for this specific call
response = await client.generate_response(
    messages=messages,
    max_tokens=16384
)

from google.genai import types

client = GeminiClient(
    config=LLMConfig(
        api_key="your-key",
        model="gemini-2.5-pro"
    ),
    thinking_config=types.ThinkingConfig(
        # Configure thinking depth and approach
    )
)

class ExtractedData(BaseModel):
    """Schema for extraction"""
    entities: list[str]
    relationships: list[dict[str, str]]

# The client automatically:
# 1. Converts Pydantic model to JSON schema
# 2. Sets response_mime_type to 'application/json'
# 3. Validates response against schema

response = await client.generate_response(
    messages=messages,
    response_model=ExtractedData
)

generation_config = types.GenerateContentConfig(
    temperature=self.temperature,
    max_output_tokens=resolved_max_tokens,
    response_mime_type='application/json',
    response_schema=ExtractedData,  # Pydantic model
    system_instruction=system_prompt,
    thinking_config=self.thinking_config
)

client = GeminiClient(
    config=LLMConfig(
        model="gemini-3-flash-preview",        # Medium
        small_model="gemini-2.5-flash-lite"    # Small
    )
)

# Uses gemini-2.5-flash-lite
response = await client.generate_response(
    messages=messages,
    model_size=ModelSize.small
)

# Uses gemini-3-flash-preview
response = await client.generate_response(
    messages=messages,
    model_size=ModelSize.medium
)

try:
    response = await client.generate_response(messages=messages)
except Exception as e:
    if 'safety' in str(e).lower() or 'blocked' in str(e).lower():
        print(f"Content blocked by safety filters: {e}")
        # No retry - content was blocked

# Example safety block details:
# "Response blocked by Gemini safety filters: HARM_CATEGORY_HARASSMENT: HIGH"

from graphiti_core.llm_client.errors import RateLimitError

try:
    response = await client.generate_response(messages=messages)
except RateLimitError as e:
    print(f"Rate limited: {e}")
    # No automatic retry - implement backoff

# On validation error:
error_context = (
    f'The previous response attempt was invalid. '
    f'Error type: {e.__class__.__name__}. '
    f'Error details: {str(e)}. '
    f'Please try again with a valid response.'
)
messages.append(Message(role='user', content=error_context))

# Looks for last valid closing bracket
array_match = re.search(r'\]\s*$', raw_output)
if array_match:
    return json.loads(raw_output[:array_match.end()])

obj_match = re.search(r'\}\s*$', raw_output)  
if obj_match:
    return json.loads(raw_output[:obj_match.end()])

client = GeminiClient()

response = await client.generate_response(
    messages=messages,
    prompt_name="summarization"
)

# Token usage from response.usage_metadata
usage = client.token_tracker.get_usage()
print(f"Prompt tokens: {usage['input_tokens']}")
print(f"Candidate tokens: {usage['output_tokens']}")
print(f"Total: {usage['total_tokens']}")

messages = [
    Message(role="system", content="You are a data extraction assistant."),
    Message(role="user", content="Extract entities from: ...")
]

# System instruction includes:
# 1. Original system message
# 2. JSON schema output instructions (if response_model provided)
# 3. Formatting guidelines

system_prompt = (
    "You are a data extraction assistant.\n\n"
    "Output ONLY valid JSON matching this schema: {...}.\n"
    "Do not include any explanatory text before or after the JSON."
)

from graphiti_core.llm_client import GeminiClient
from graphiti_core.llm_client.config import LLMConfig, ModelSize
from pydantic import BaseModel

class Classification(BaseModel):
    category: str
    confidence: float

client = GeminiClient(
    config=LLMConfig(
        model="gemini-3-flash-preview",
        small_model="gemini-2.5-flash-lite"
    )
)

items = ["text 1", "text 2", "text 3"]
results = []

for item in items:
    messages = [
        Message(role="system", content="Classify the text."),
        Message(role="user", content=item)
    ]
    
    result = await client.generate_response(
        messages=messages,
        response_model=Classification,
        model_size=ModelSize.small  # Use faster model
    )
    results.append(result)

# The client checks prompt_feedback.block_reason
if prompt_feedback and block_reason:
    raise Exception(f'Prompt blocked by Gemini: {block_reason}')