from typing import Optional, List, Tuple, Any
from google import genai
import os
import instructor
def setup_gemini_client(
    user_id: Optional[str] = None,
    session_id: Optional[str] = None,
    session_name: Optional[str] = None,
    session_path: Optional[str] = None,
    custom_properties: Optional[dict] = None,
    other_metadata: Optional[dict] = None
) -> Any:
    """
    Factory function to create a Gemini client with user-specific Helicone headers.
    Args:
        user_id: Optional user identifier for Helicone tracking
        session_id: Optional session identifier for Helicone tracking
        session_name: Optional name for the session in Helicone
        session_path: Optional path for the session in Helicone
        custom_properties: Optional dictionary of custom properties for Helicone
        other_metadata: Optional dictionary of additional metadata headers
    Returns:
        Configured Gemini client with Instructor integration
    """
    # Base headers required for Helicone
    metadata: List[Tuple[str, str]] = [
        ("helicone-auth", f"Bearer {os.environ.get('HELICONE_API_KEY')}"),
        ("helicone-target-url", "https://generativelanguage.googleapis.com"),
    ]
    # Add user_id if provided
    if user_id:
        metadata.append(("helicone-user-id", user_id))
    # Add session_id if provided
    if session_id:
        metadata.append(("helicone-session-id", session_id))
    # Add session_name if provided
    if session_name:
        metadata.append(("helicone-session-name", session_name))
    # Add session_path if provided
    if session_path:
        metadata.append(("helicone-session-path", session_path))
    # Add custom properties if provided
    if custom_properties:
        for key, value in custom_properties.items():
            metadata.append((f"helicone-property-{key}", value))
    # Add other metadata if provided
    if other_metadata:
        for key, value in other_metadata.items():
            metadata.append((key, value))
    # Configure the client with metadata
    genai.configure(
        api_key=os.environ.get('GOOGLE_API_KEY'),
        client_options={
            "api_endpoint": "gateway.helicone.ai",
        },
        default_metadata=metadata,
        transport="rest",
    )
    # Create Gemini client with Instructor integration
    gemini_client = instructor.from_gemini(
        genai.GenerativeModel(
            model_name="gemini-2.0-flash",
        ),
        mode=instructor.Mode.GEMINI_JSON,
    )
    return gemini_client
session_client = setup_gemini_client(
    user_id="user123",
    session_id="session456",
    session_name="Customer Support Chat",
    session_path="/support/tickets/123",
    custom_properties={
        "job_id": "1234567890",
        "job_name": "Customer Support Chat",
    }
)
response = session_client.chat.completions.create(
    messages=[
        {"role": "user", "content": "What are black holes?"}
    ],
    response_model=Response
)