Vedanshipanda/layer10-api

1import os2import json3import time4import google.generativeai as genai5from dotenv import load_dotenv6 7load_dotenv()8 9# Setup Gemini10api_key = os.getenv("GEMINI_API_KEY")11if not api_key:12    raise ValueError("GEMINI_API_KEY not found in .env")13 14genai.configure(api_key=api_key)15model = genai.GenerativeModel('models/gemini-flash-latest',16                              generation_config={"response_mime_type": "application/json"})17 18RAW_DATA_PATH = os.path.join("data", "raw_corpus.json")19OUTPUT_PATH = os.path.join("data", "extracted_graph.json")20 21def load_data():22    if not os.path.exists(RAW_DATA_PATH):23        print(f"❌ Error: {RAW_DATA_PATH} not found.")24        return [], []25    with open(RAW_DATA_PATH, "r", encoding="utf-8") as f:26        raw_data = json.load(f)27    28    extracted_data = []29    if os.path.exists(OUTPUT_PATH):30        try:31            with open(OUTPUT_PATH, "r", encoding="utf-8") as f:32                extracted_data = json.load(f)33        except: pass34    return raw_data, extracted_data35 36def save_progress(data):37    with open(OUTPUT_PATH, "w", encoding="utf-8") as f:38        json.dump(data, f, indent=4)39    print(f"💾 Saved {len(data)} items to disk.")40 41def extract_batch(items):42    batched_text = ""43    for i, item in enumerate(items):44        uid = item.get("html_url", str(item.get("id")))45        title = item.get("title") or ""46        # FIX: Ensure body is a string even if None, before slicing47        body = item.get("body") or "" 48        49        batched_text += f"\n--- ITEM {i} (ID: {uid}) ---\nTITLE: {title}\nBODY: {body[:1500]}\n"50 51    prompt = f"""52    You are a Knowledge Graph extractor. Analyze these {len(items)} GitHub items.53    54    For EACH item, extract:55    1. Entities: "Person", "Feature", "Bug", "Artifact", "Topic", "Decision".56    2. Relationships: How entities interact (e.g., "Reported", "Fixed", "Affects").57    58    CRITICAL: 59    - Use EXACT entity names from the text.60    - 'source' and 'target' in relationships must match Entity names exactly.61    - Provide a short 'text_excerpt' as evidence for every entity.62 63    Output a JSON LIST of objects:64    [65      {{66        "source_id": "ID_FROM_HEADER",67        "graph_data": {{68          "entities": [{{"name": "...", "type": "...", "text_excerpt": "..."}}],69          "relationships": [{{"source": "...", "target": "...", "relation_type": "..."}}]70        }}71      }}72    ]73 74    Data to process:75    {batched_text}76    """77    78    try:79        response = model.generate_content(prompt)80        return json.loads(response.text)81    except Exception as e:82        print(f"❌ API Error: {e}")83        return None84 85def process_corpus():86    raw_data, extracted_data = load_data()87    processed_ids = {item["source_id"] for item in extracted_data}88    remaining = [x for x in raw_data if x.get("html_url", str(x.get("id"))) not in processed_ids]89 90    print(f"🚀 Starting Extraction with gemini-flash-latest. Remaining items: {len(remaining)}")91 92    batch_size = 593    for i in range(0, len(remaining), batch_size):94        batch = remaining[i : i + batch_size]95        print(f"📦 Processing batch {i//batch_size + 1} ({len(batch)} items)...")96        97        results = extract_batch(batch)98        if results:99            extracted_data.extend(results)100            save_progress(extracted_data)101        102        # 75s wait ensures we NEVER hit the quota window limit103        if i + batch_size < len(remaining):104            print("⏳ Cooldown 75s to reset API quota...")105            time.sleep(75)106 107    print("\n✅ Extraction Complete!")108 109if __name__ == "__main__":110    process_corpus()