Gemini for Data Analytics

studio_ref = gda.StudioDatasourceReference(
    datasource_id="your-studio-datasource-id"
)

datasource_refs = gda.DatasourceReferences(
    studio=gda.StudioDatasourceReferences(
        studio_references=[studio_ref]
    )
)

system_instruction = "You are a sales analyst for an e-commerce company. Always include percentages and comparisons to previous periods."

context = gda.Context(
    system_instruction=system_instruction,
    datasource_references=datasource_refs,
)

system_instruction = """
- system_instruction: >
    You are an expert sales analyst for Acme Corp e-commerce store.
    
- tables:
    - table:
        - name: acme.sales.orders
        - description: Orders from the e-commerce store
        - synonyms: [sales, transactions]
        - tags: [sale, order, revenue]
        - fields:
            - field:
                - name: order_id
                - description: Unique identifier for each order
            - field:
                - name: customer_id
                - description: Unique customer identifier
            - field:
                - name: status
                - description: Order status
                - sample_values: [complete, shipped, returned, pending]
            - field:
                - name: total_amount
                - description: Total order value in USD
                - aggregations: [sum, avg, max, min]
            - field:
                - name: created_at
                - description: Order creation timestamp
        - measures:
            - measure:
                - name: revenue
                - description: Total revenue (sum of total_amount)
                - exp: SUM(total_amount)
        - golden_queries:
            - golden_query:
                - natural_language_query: What were total sales last month?
                - sql_query: >
                    SELECT SUM(total_amount) as revenue
                    FROM acme.sales.orders
                    WHERE created_at >= DATE_SUB(CURRENT_DATE(), INTERVAL 1 MONTH)
                    AND status = 'complete'
                    
- glossaries:
    - glossary:
        - term: Revenue
        - description: Total value of completed orders
        - synonyms: [sales, income, earnings]
    - glossary:
        - term: Conversion Rate
        - description: Percentage of visitors who make a purchase
        - synonyms: [CVR, conversion]
        
- additional_descriptions:
    - text: All monetary values are in USD unless specified otherwise.
    - text: Fiscal year starts in April.
"""

chat_client = gda.DataChatServiceClient()

# Create conversation
conversation = chat_client.create_conversation(
    parent=f"projects/{PROJECT_ID}/locations/global",
    conversation_id="sales_analysis_1",
    conversation=gda.Conversation(
        agents=[agent.name]
    ),
)

# First question
response1 = chat_client.generate_message(
    conversation=conversation.name,
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="Show me sales by region for Q4 2024"
        )
    )],
)

# Follow-up question (context maintained)
response2 = chat_client.generate_message(
    conversation=conversation.name,
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="Now compare that to Q4 2023"  # "that" refers to previous query
        )
    )],
)

# Another follow-up
response3 = chat_client.generate_message(
    conversation=conversation.name,
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="Show it as a bar chart"  # "it" refers to the comparison
        )
    )],
)

def handle_response(response):
    for msg in response.messages:
        system_msg = msg.system_message
        
        # Text response
        if system_msg.text:
            print("Text:", system_msg.text.parts)
        
        # Schema information
        elif system_msg.schema:
            print("Schema resolved:")
            for datasource in system_msg.schema.result.datasources:
                print(f"  Table: {datasource.bigquery_table_reference.table_id}")
        
        # Data query results
        elif system_msg.data:
            if system_msg.data.generated_sql:
                print("SQL:", system_msg.data.generated_sql)
            
            if system_msg.data.result:
                # Convert to pandas DataFrame
                import pandas as pd
                
                fields = [f.name for f in system_msg.data.result.schema.fields]
                rows = [{f: row[f] for f in fields} for row in system_msg.data.result.data]
                df = pd.DataFrame(rows)
                print(df)
        
        # Chart visualization
        elif system_msg.chart:
            import altair as alt
            import json
            
            # Render Vega-Lite chart
            vega_config = system_msg.chart.result.vega_config
            chart = alt.Chart.from_dict(dict(vega_config))
            chart.display()
        
        # Clarification needed
        elif system_msg.clarification:
            print("Please clarify:")
            for question in system_msg.clarification.questions:
                print(f"  {question.question}")
                for option in question.options:
                    print(f"    - {option}")

handle_response(response)

stream = chat_client.generate_message_stream(
    conversation=conversation.name,
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="Analyze sales trends over the past year"
        )
    )],
)

for chunk in stream:
    for msg in chunk.messages:
        if msg.system_message.text:
            print(msg.system_message.text.parts, end="", flush=True)

chat_client = gda.DataChatServiceClient()

# Define context inline
inline_context = gda.Context(
    system_instruction="You are a data analyst",
    datasource_references=datasource_refs,
)

# Generate response without creating agent or conversation
response = chat_client.generate_message(
    inline_context=inline_context,
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="How many airports are in Texas?"
        )
    )],
)

context = gda.Context(
    system_instruction="You are a data scientist",
    datasource_references=datasource_refs,
    options=gda.ConversationOptions(
        analysis=gda.AnalysisOptions(
            python=gda.AnalysisOptions.Python(
                enabled=True  # Enable Python for complex calculations
            )
        )
    ),
)

# Agent can now use Python for:
# - Statistical analysis
# - Data transformations
# - Complex calculations
# - Custom visualizations

response = chat_client.generate_message(
    conversation=conversation.name,
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="Create a line chart of monthly sales with a trend line"
        )
    )],
    visualization_config=gda.VisualizationConfig(
        chart_type="line",
        show_legend=True,
        theme="dark",
    ),
)

# Combine BigQuery and Looker
bq_ref = gda.BigQueryTableReference(
    project_id="my-project",
    dataset_id="warehouse",
    table_id="raw_events",
)

looker_ref = gda.LookerExploreReference(
    looker_instance_uri="https://company.looker.com",
    lookml_model="business_intelligence",
    explore="sales_metrics",
)

datasource_refs = gda.DatasourceReferences(
    bq=gda.BigQueryTableReferences(table_references=[bq_ref]),
    looker=gda.LookerExploreReferences(explore_references=[looker_ref]),
)

agent = client.create_data_agent(
    parent=f"projects/{PROJECT_ID}/locations/global",
    data_agent=gda.DataAgent(
        data_analytics_agent=gda.DataAnalyticsAgent(
            published_context=gda.Context(
                system_instruction="""You have access to both raw event data (BigQuery)
                and business metrics (Looker). Use raw data for detailed analysis
                and Looker for standard business metrics.""",
                datasource_references=datasource_refs,
            )
        )
    ),
)

# Query spans both sources
response = chat_client.generate_message(
    inline_context=gda.Context(
        datasource_references=datasource_refs,
    ),
    messages=[gda.Message(
        user_message=gda.UserMessage(
            text="Compare funnel conversion rates from raw events with the aggregated metrics in Looker"
        )
    )],
)

# List conversations
for conv in chat_client.list_conversations(
    parent=f"projects/{PROJECT_ID}/locations/global"
):
    print(f"Conversation: {conv.name}")

# Get conversation history
conv = chat_client.get_conversation(
    name="projects/my-project/locations/global/conversations/conv_123"
)

for msg in conv.messages:
    if msg.user_message:
        print(f"User: {msg.user_message.text}")
    elif msg.system_message:
        print(f"Agent: {msg.system_message.text.parts}")

# Delete conversation
chat_client.delete_conversation(
    name="projects/my-project/locations/global/conversations/conv_123"
)