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/review."""
    llm_config = config.get("llm_pipeline", 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)

    # Detect Ollama native API (11434 port) — use /api/chat instead of /v1/chat/completions
    is_ollama_native = ":11434" in base_url

    if is_ollama_native:
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": False,
            "options": {"num_predict": max_tokens, "temperature": 0.1},
        }
        try:
            resp = requests.post(f"{base_url}/api/chat", json=payload, headers={"Content-Type": "application/json"}, timeout=120)
            if resp.status_code == 200:
                return resp.json().get("message", {}).get("content", "")
            else:
                return f"LLM error: {resp.status_code}"
        except Exception as e:
            return f"LLM error: {str(e)}"
    else:
        # OpenAI-compatible format
        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()
                msg = data["choices"][0]["message"]
                content = (msg.get("content") or msg.get("reasoning_content") or "").strip()
                if not content and msg.get("reasoning_content"):
                    rc = msg["reasoning_content"]
                    import re
                    json_match = re.search(r'(\{.*\})', rc, re.DOTALL)
                    if json_match:
                        content = json_match.group()
                return 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. Analyze a GitHub repository and determine if it contains a reusable AI workflow or pattern that another agent could learn from and actually implement.

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 with type description", "Input 2 with type description"],
  "steps": [
    "Step 1: Describe the specific action, mentioning the exact tool/function/file used (e.g. 'Run langgraph chain with agent.py')",
    "Step 2: ...",
    "Step 3: ..."
  ],
  "outputs": ["Output 1 with description", "Output 2 with description"],
  "failure_modes": ["Specific failure scenario with mitigation"],
  "confidence": 0.95,
  "reusable": true,
  "general_purpose": false,
  "explanation": "Why this is reusable",
  "implementation_details": {{
    "framework": "e.g. langchain, langgraph, autogen, crewai, custom",
    "dependencies": ["python-packages-needed"],
    "key_files": ["path/to/key_file.py - description"],
    "code_snippets": ["Brief but concrete code or config example from the repo"],
    "setup_steps": ["Prerequisite setup commands or configs"]
  }}
}}

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

CRITICAL: Steps must be SPECIFIC — mention actual file names, function calls, tool names, or configuration details from the repository. A step like 'Researcher agent gathers facts' is too vague. Instead: 'Researcher agent (agent.py) uses LangGraph create_react_agent with SerperDevTool to gather facts.'

Criteria for a reusable workflow:
- It describes a concrete process, not just a tool or library
- Steps mention specific implementations from the code
- It has clear inputs, steps, and outputs
- It could be adapted to different contexts
- It has at least 3 distinct steps
- It solves a real problem

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,
    }