lakehouse/scripts/extract_knowledge.py

#!/usr/bin/env python3
"""
Knowledge extraction pipeline: process all conversations through LLM,
extract structured findings, build knowledge graph as Parquet tables.
"""

import json, urllib.request, time, re

API = "http://localhost:3100"

def sql(query):
    req = urllib.request.Request(f"{API}/query/sql",
        data=json.dumps({"sql": query}).encode(),
        headers={"Content-Type": "application/json"}, method="POST")
    return json.loads(urllib.request.urlopen(req, timeout=30).read())

def generate(prompt, max_tokens=512):
    req = urllib.request.Request(f"{API}/ai/generate",
        data=json.dumps({"prompt": prompt, "max_tokens": max_tokens, "temperature": 0.1}).encode(),
        headers={"Content-Type": "application/json"}, method="POST")
    return json.loads(urllib.request.urlopen(req, timeout=120).read())

def ingest_csv(name, csv_text):
    boundary = "----kb" + str(int(time.time()))
    body = f"--{boundary}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"{name}.csv\"\r\nContent-Type: text/csv\r\n\r\n{csv_text}\r\n--{boundary}--\r\n"
    req = urllib.request.Request(f"{API}/ingest/file?name={name}",
        data=body.encode(), headers={"Content-Type": f"multipart/form-data; boundary={boundary}"}, method="POST")
    try: return urllib.request.urlopen(req, timeout=15).read().decode()
    except: return "error"

# Get all conversation datasets
old = ['call_log','placements','clients','threat_intel','timesheets','meta_runs','candidates',
       'job_orders','lab_trials','users','email_log','new_candidates','schema_test_v1',
       'schema_test_v2','stress_test_ingest','stress_ingest','knowledge_topics',
       'knowledge_decisions','knowledge_insights']
ds = json.loads(urllib.request.urlopen(f"{API}/catalog/datasets").read())
convs = [d for d in ds if d['name'] not in old]

print(f"Processing {len(convs)} conversations...")

# Knowledge tables
all_topics = []      # (conversation, topic)
all_decisions = []   # (conversation, decision, context)
all_insights = []    # (conversation, page, insight, reasoning)
all_questions = []   # (conversation, question)

for conv in convs:
    name = conv['name']
    pages = conv.get('row_count', 0) or 0
    if pages == 0: continue

    print(f"\n{'='*50}")
    print(f"Extracting from: {name} ({pages} pages)")

    # Sample pages (every 5th page for large convos, all for small)
    step = max(1, pages // 20)  # ~20 samples max per conversation
    r = sql(f"SELECT turn_number, content FROM {name} WHERE turn_number % {step} = 1 ORDER BY turn_number LIMIT 20")

    for row in r['rows']:
        page_num = row.get('turn_number', 0)
        content = (row.get('content', '') or '')[:2000]
        if len(content) < 100: continue

        prompt = f"""Analyze this excerpt from a technical AI conversation. Extract structured knowledge.

TEXT:
{content}

Output a JSON object with:
- topics: list of technology/concept topics (max 5)
- decisions: list of decisions made or recommended (max 3)
- key_insight: single most important takeaway (1 sentence)
- open_questions: unresolved questions (max 2)

Output ONLY valid JSON."""

        try:
            resp = generate(prompt, 300)
            text = resp['text'].strip()
            # Extract JSON from response
            if '```' in text:
                text = text.split('```')[1].replace('json','',1).strip()
            # Find JSON object
            start = text.find('{')
            end = text.rfind('}') + 1
            if start >= 0 and end > start:
                data = json.loads(text[start:end])

                for topic in data.get('topics', [])[:5]:
                    all_topics.append(f"{name},{page_num},{topic}")

                for decision in data.get('decisions', [])[:3]:
                    all_decisions.append(f"{name},{page_num},{decision}")

                insight = data.get('key_insight', '')
                if insight:
                    all_insights.append(f"{name},{page_num},{insight}")

                for q in data.get('open_questions', [])[:2]:
                    all_questions.append(f"{name},{page_num},{q}")

                print(f"  p{page_num}: {len(data.get('topics',[]))} topics, {len(data.get('decisions',[]))} decisions")
            else:
                print(f"  p{page_num}: no valid JSON")
        except Exception as e:
            print(f"  p{page_num}: error — {str(e)[:50]}")

# Build and ingest knowledge tables
print(f"\n{'='*50}")
print(f"Extracted: {len(all_topics)} topics, {len(all_decisions)} decisions, {len(all_insights)} insights, {len(all_questions)} questions")

if all_topics:
    csv = "conversation,page,topic\n" + "\n".join(all_topics)
    print(f"Ingesting knowledge_topics ({len(all_topics)} rows)...", end=" ")
    print(ingest_csv("knowledge_topics", csv))

if all_decisions:
    csv = "conversation,page,decision\n" + "\n".join(all_decisions)
    print(f"Ingesting knowledge_decisions ({len(all_decisions)} rows)...", end=" ")
    print(ingest_csv("knowledge_decisions", csv))

if all_insights:
    csv = "conversation,page,insight\n" + "\n".join(all_insights)
    print(f"Ingesting knowledge_insights ({len(all_insights)} rows)...", end=" ")
    print(ingest_csv("knowledge_insights", csv))

if all_questions:
    csv = "conversation,page,question\n" + "\n".join(all_questions)
    print(f"Ingesting knowledge_questions ({len(all_questions)} rows)...", end=" ")
    print(ingest_csv("knowledge_questions", csv))

print(f"\nDone. Query with:")
print(f"  SELECT topic, COUNT(*) FROM knowledge_topics GROUP BY topic ORDER BY COUNT(*) DESC")
print(f"  SELECT * FROM knowledge_decisions")
print(f"  SELECT * FROM knowledge_insights")