agent-skills/pipeline/extractor.py

"""Stage 4: Workflow Extractor — Extract reusable workflow from repo context."""
import json
import os
import requests
import re


def call_llm(prompt, config):
    """Call the configured LLM for extraction."""
    llm_config = config.get("llm", {})
    base_url = llm_config.get("base_url", "http://100.64.0.2:8083/v1")
    model = llm_config.get("model", "")
    api_key = llm_config.get("api_key", "")
    max_tokens = llm_config.get("max_tokens", 8000)

    headers = {
        "Content-Type": "application/json",
    }
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": prompt},
        ],
        "max_tokens": max_tokens,
        "temperature": 0.1,
    }

    try:
        resp = requests.post(f"{base_url}/v1/chat/completions", json=payload, headers=headers, timeout=120)
        if resp.status_code == 200:
            data = resp.json()
            return data["choices"][0]["message"]["content"]
        else:
            return f"LLM error: {resp.status_code} {resp.text[:200]}"
    except Exception as e:
        return f"LLM error: {str(e)}"


def extract_workflow(reader_output, config):
    """
    Attempt to extract a reusable workflow from the repo context.
    Returns structured workflow or rejection.
    """
    repo = reader_output.get("repository", "")
    content_sections = reader_output.get("content", {})

    # Build context for LLM
    context_parts = []
    for path, content in content_sections.items():
        context_parts.append(f"--- {path} ---\n{content[:15000]}")

    context = "\n\n".join(context_parts)

    prompt = f"""You are a workflow extractor. Your job is to analyze a GitHub repository and determine if it contains a reusable AI workflow or pattern that another agent could learn from.

If the repository contains a reusable workflow, extract it into this exact JSON structure:
{{
  "has_workflow": true,
  "skill_name": "short-descriptive-name",
  "goal": "One sentence: what this workflow accomplishes",
  "inputs": ["Input 1", "Input 2"],
  "steps": ["Step 1", "Step 2", "Step 3"],
  "outputs": ["Output 1", "Output 2"],
  "failure_modes": ["What can go wrong"],
  "confidence": 0.95,
  "reusable": true,
  "general_purpose": true,
  "explanation": "Why this is reusable and general-purpose"
}}

If the repository does NOT contain a reusable workflow, return:
{{
  "has_workflow": false,
  "reason": "Why no reusable workflow was found"
}}

Criteria for a reusable workflow:
- It describes a process or pattern, not just a tool or library
- It has clear inputs, steps, and outputs
- It could be applied to different contexts outside this specific repo
- It has at least 3 distinct steps
- It solves a real problem, not a toy example

Repository: {repo}

Repository context:
{context[:30000]}

Return ONLY valid JSON. No markdown, no explanation outside the JSON."""

    result_text = call_llm(prompt, config)

    # Parse JSON from response — try multiple strategies
    workflow = None
    for attempt in range(3):
        try:
            # Strip markdown code blocks if present
            cleaned = result_text.strip()
            if cleaned.startswith("```"):
                cleaned = cleaned.split("```")[1]
                if cleaned.startswith("json"):
                    cleaned = cleaned[4:]
                cleaned = cleaned.rstrip("```")
            cleaned = cleaned.strip()

            # Try to find JSON object in text
            json_match = re.search(r'\{.*\}', cleaned, re.DOTALL)
            if json_match:
                cleaned = json_match.group()

            workflow = json.loads(cleaned)
            break
        except json.JSONDecodeError:
            if attempt < 2:
                # Retry with a simpler prompt
                prompt = f"""Extract a reusable workflow from this repo as JSON. Return ONLY valid JSON.

Repo: {repo}
Context (first 10000 chars): {context[:10000]}

Format:
{{"has_workflow": true/false, "skill_name": "...", "goal": "...", "inputs": [...], "steps": [...], "outputs": [...], "failure_modes": [...], "confidence": 0-1, "reusable": true/false, "general_purpose": true/false, "explanation": "..."}}

If no reusable workflow: {{"has_workflow": false, "reason": "..."}}

Return ONLY JSON. No markdown."""
                result_text = call_llm(prompt, config)
            else:
                break

    if workflow is None:
        return {
            "status": "PARSE_ERROR",
            "raw": result_text[:500],
            "repository": repo,
        }

    if not workflow.get("has_workflow", False):
        return {
            "status": "NO_WORKFLOW",
            "reason": workflow.get("reason", "Extractor determined no reusable workflow"),
            "repository": repo,
        }

    # Validate minimum requirements
    steps = workflow.get("steps", [])
    if len(steps) < 3:
        return {
            "status": "NO_WORKFLOW",
            "reason": f"Only {len(steps)} steps found, minimum is 3",
            "repository": repo,
        }

    if workflow.get("confidence", 0) < 0.7:
        return {
            "status": "LOW_CONFIDENCE",
            "confidence": workflow.get("confidence"),
            "repository": repo,
        }

    return {
        "status": "EXTRACTED",
        "repository": repo,
        "workflow": workflow,
    }