"""LLM client for generating documentation. Supports Ollama and OpenAI-compatible APIs.""" import requests import os import concurrent.futures import time from pathlib import Path # Load .env if present _env_file = Path(__file__).parent / ".env" if _env_file.exists(): for line in _env_file.read_text().splitlines(): line = line.strip() if line and not line.startswith("#") and "=" in line: key, _, val = line.partition("=") os.environ.setdefault(key.strip(), val.strip()) # Backend: "ollama" or "openai" LLM_BACKEND = os.environ.get("LLM_BACKEND", "ollama") # Ollama settings OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://192.168.86.172:11434") OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen2.5:7b") # OpenAI-compatible settings (works with Kiro gateway, OpenRouter, etc.) OPENAI_URL = os.environ.get("OPENAI_URL", "http://192.168.86.11:8000") OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "claude-haiku-4") OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "not-needed") MAX_CONCURRENT = int(os.environ.get("MAX_CONCURRENT", "4")) def generate_file_doc(filepath: str, content: str) -> str: """Generate documentation for a single file.""" if len(content) > 8000: content = content[:8000] + "\n\n... [truncated]" prompt = f"""You are a senior software engineer documenting a Go codebase. Describe what this file does in 2-4 sentences. Be specific about: - The domain logic and purpose (not just "this file contains functions") - Key types, interfaces, or structs defined - How it fits into the larger system (if apparent from imports/naming) Do NOT describe Go syntax or language mechanics. Describe WHAT the code does and WHY. File: {filepath} ```go {content} ``` Documentation:""" return _call_llm(prompt) def generate_relationship_doc(file_a: str, content_a: str, file_b: str, content_b: str) -> str: """Generate documentation for a relationship between two files.""" if len(content_a) > 4000: content_a = content_a[:4000] + "\n... [truncated]" if len(content_b) > 4000: content_b = content_b[:4000] + "\n... [truncated]" prompt = f"""You are a senior software engineer documenting how two files in a Go codebase interact. Describe in 1-2 sentences how File A uses or depends on File B. Be specific about which types, functions, or interfaces are shared. File A: {file_a} ```go {content_a} ``` File B: {file_b} ```go {content_b} ``` Relationship:""" return _call_llm(prompt) def generate_repo_doc(readme: str, entry_files: list[tuple[str, str]]) -> str: """Generate repo-level documentation from README and key entry points.""" files_section = "" for path, content in entry_files[:5]: snippet = content[:2000] if len(content) > 2000 else content files_section += f"\n--- {path} ---\n{snippet}\n" readme_section = readme[:3000] if len(readme) > 3000 else readme prompt = f"""You are a senior software engineer writing a project overview. Based on the README and key source files below, write a 4-6 sentence summary of this project. Cover: - What the project does (its purpose) - Key architectural patterns (routing, middleware, etc.) - The main abstractions and how they fit together README: {readme_section} Key source files: {files_section} Project Overview:""" return _call_llm(prompt) def generate_docs_batch(items: list[tuple[str, str]], doc_fn) -> list[str]: """Generate docs for multiple items concurrently.""" results = [None] * len(items) with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor: future_to_idx = {} for i, (filepath, content) in enumerate(items): future = executor.submit(doc_fn, filepath, content) future_to_idx[future] = i done = 0 total = len(items) for future in concurrent.futures.as_completed(future_to_idx): idx = future_to_idx[future] try: results[idx] = future.result() except Exception as e: results[idx] = f"[doc generation failed: {e}]" done += 1 if done % 10 == 0 or done == total: print(f" Generated {done}/{total} docs") return results def _call_ollama(prompt: str, retries: int = 3) -> str: """Call Ollama API with retries.""" for attempt in range(retries): try: resp = requests.post( f"{OLLAMA_URL}/api/generate", json={ "model": OLLAMA_MODEL, "prompt": prompt, "stream": False, "options": { "temperature": 0.3, "num_predict": 256, }, }, timeout=120, ) resp.raise_for_status() return resp.json()["response"].strip() except Exception as e: if attempt < retries - 1: time.sleep(2 ** attempt) continue return f"[doc generation failed after {retries} attempts: {e}]" def _call_openai(prompt: str, retries: int = 3) -> str: """Call OpenAI-compatible API (Kiro gateway, OpenRouter, etc.).""" for attempt in range(retries): try: resp = requests.post( f"{OPENAI_URL}/v1/chat/completions", headers={"Authorization": f"Bearer {OPENAI_API_KEY}"}, json={ "model": OPENAI_MODEL, "messages": [{"role": "user", "content": prompt}], "temperature": 0.3, "max_tokens": 256, }, timeout=120, ) resp.raise_for_status() return resp.json()["choices"][0]["message"]["content"].strip() except Exception as e: if attempt < retries - 1: time.sleep(2 ** attempt) continue return f"[doc generation failed after {retries} attempts: {e}]" def _call_llm(prompt: str) -> str: """Route to the configured backend.""" if LLM_BACKEND == "openai": return _call_openai(prompt) return _call_ollama(prompt)