Team Ai
Apppublic

operations-granite/HuggingFace-Granite-AI-Practice-Agent-Dev

sourceHugging Faceapache-2.0updated 3mo agoView on Hugging Face
0likes
test_parallel_research.py146 linesDownload Raw Back to root
1#!/usr/bin/env python32"""3Test parallel research approach - exactly like production will use it4"""5import os6import asyncio7from dotenv import load_dotenv8from google import genai9from google.genai import types10from pydantic import BaseModel, Field11from typing import List12 13load_dotenv()14 15GEMINI_MODEL = os.getenv('GEMINI_MODEL', 'gemini-3-flash-preview')16 17# Exact schemas from production18class IntelligenceSection(BaseModel):19    intro: str20    items: List[str]21    forSales: str22 23class CompanyMetadata(BaseModel):24    industry: str25    location: str26    employee_count: str27 28# Group schemas29class MarketIntelligence(BaseModel):30    metadata: CompanyMetadata31    company_overview: IntelligenceSection32    recent_news: IntelligenceSection33    strategic_priorities: IntelligenceSection34    technology_adoption: IntelligenceSection35    budget_financial_health: IntelligenceSection36    competitive_landscape: IntelligenceSection37 38class GraniteAlignment(BaseModel):39    granite_capabilities_match: IntelligenceSection40    recommended_partners: IntelligenceSection41    capability_gaps: IntelligenceSection42 43class SalesStrategy(BaseModel):44    decision_making_process: IntelligenceSection45    past_projects_vendors: IntelligenceSection46    pain_points_challenges: IntelligenceSection47    value_proposition: IntelligenceSection48    recommended_approach: IntelligenceSection49    likely_objections: IntelligenceSection50    talking_points: List[str]51 52client = genai.Client(api_key=os.getenv('GOOGLE_API_KEY'))53grounding_tool = types.Tool(google_search=types.GoogleSearch())54 55# Configs for each group56config_market = types.GenerateContentConfig(57    tools=[grounding_tool],58    temperature=0.7,59    max_output_tokens=65536,60    response_mime_type="application/json",61    response_schema=MarketIntelligence,62    thinking_config=types.ThinkingConfig(thinking_budget=0)63)64 65config_granite = types.GenerateContentConfig(66    tools=[grounding_tool],67    temperature=0.7,68    max_output_tokens=65536,69    response_mime_type="application/json",70    response_schema=GraniteAlignment,71    thinking_config=types.ThinkingConfig(thinking_budget=0)72)73 74config_sales = types.GenerateContentConfig(75    tools=[grounding_tool],76    temperature=0.7,77    max_output_tokens=65536,78    response_mime_type="application/json",79    response_schema=SalesStrategy,80    thinking_config=types.ThinkingConfig(thinking_budget=0)81)82 83prompt = """84Research Saudi Telecom Company (STC).85 86Each section: 250-350 words with specific names, dates, dollar amounts.87Use Google Search for current 2025-2026 information.88"""89 90async def main():91    print("Testing PARALLEL research approach")92    print("=" * 60)93 94    try:95        # Create parallel tasks96        print("Launching 3 parallel API calls...")97        market_task = asyncio.create_task(asyncio.to_thread(98            client.models.generate_content,99            model=GEMINI_MODEL,100            contents=prompt,101            config=config_market102        ))103 104        granite_task = asyncio.create_task(asyncio.to_thread(105            client.models.generate_content,106            model=GEMINI_MODEL,107            contents=prompt,108            config=config_granite109        ))110 111        sales_task = asyncio.create_task(asyncio.to_thread(112            client.models.generate_content,113            model=GEMINI_MODEL,114            contents=prompt,115            config=config_sales116        ))117 118        # Wait for all to complete119        print("Waiting for parallel completion...")120        market_resp, granite_resp, sales_resp = await asyncio.gather(121            market_task, granite_task, sales_task122        )123 124        print(f"\n✅ ALL 3 CALLS SUCCEEDED!")125        print(f"   Market Intelligence: {len(market_resp.text)} chars")126        print(f"   Granite Alignment: {len(granite_resp.text)} chars")127        print(f"   Sales Strategy: {len(sales_resp.text)} chars")128        print(f"   Total: {len(market_resp.text) + len(granite_resp.text) + len(sales_resp.text)} chars")129 130        # Parse and combine131        import json132        market = json.loads(market_resp.text)133        granite = json.loads(granite_resp.text)134        sales = json.loads(sales_resp.text)135 136        combined = {**market, **granite, **sales}137        print(f"\n   Combined sections: {len(combined)} fields")138        print(f"   Has metadata: {'metadata' in combined}")139        print(f"   Has talking_points: {'talking_points' in combined}")140 141    except Exception as e:142        print(f"\n❌ FAILED: {e}")143 144if __name__ == "__main__":145    asyncio.run(main())146