netdocai 0.1.0__py3-none-any.whl

netdocai/agent.py

@@ -0,0 +1,320 @@
+import os
+import sys
+from typing import Any, TypedDict
+
+from google.genai import types
+from .gemini import generate_content
+from langgraph.graph import END, START, StateGraph
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from .prompts import SYSTEM_PROMPT, build_user_prompt
+from .report import build_report
+
+
+MAX_STEPS = 8
+
+
+class AgentState(TypedDict):
+    problem: str
+    history: list[dict[str, Any]]
+    tool_definitions: list[dict[str, Any]]
+    pending_calls: list[dict[str, Any]]
+    observations: list[dict[str, Any]]
+    hypothesis: str
+    response: str
+    steps: int
+
+
+def build_mcp_tools(
+    tool_definitions: list[dict[str, Any]],
+) -> list[types.Tool]:
+    declarations = []
+
+    for tool in tool_definitions:
+        declarations.append(
+            types.FunctionDeclaration(
+                name=tool["name"],
+                description=tool["description"],
+                parameters_json_schema=tool["input_schema"],
+            )
+        )
+
+    return [types.Tool(function_declarations=declarations)]
+
+
+async def get_tool_definitions() -> list[dict[str, Any]]:
+    server_params = StdioServerParameters(
+        command=sys.executable,
+        args=["-m", "netdocai.mcp_server"],
+        env=os.environ.copy(),
+    )
+
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write) as session:
+            await session.initialize()
+
+            result = await session.list_tools()
+
+            return [
+                {
+                    "name": tool.name,
+                    "description": tool.description or "",
+                    "input_schema": tool.input_schema,
+                }
+                for tool in result.tools
+            ]
+
+
+async def plan(state: AgentState) -> AgentState:
+    contents = [
+        types.Content.model_validate(message)
+        for message in state["history"]
+    ]
+
+    gemini_tools = build_mcp_tools(state["tool_definitions"])
+
+    response = await generate_content(
+        contents=contents,
+        config=types.GenerateContentConfig(
+            tools=gemini_tools,
+            automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                disable=True
+            ),
+        ),
+    )
+
+    function_calls = response.function_calls
+
+    if not function_calls:
+        print("[Agent] Finalizing diagnosis...")
+
+        report = build_report(
+            problem=state["problem"],
+            observations=state["observations"],
+            hypothesis=state["hypothesis"],
+            conclusion=response.text or "",
+        )
+
+        return {
+            **state,
+            "response": report,
+            "pending_calls": [],
+        }
+
+    model_content = response.candidates[0].content
+
+    updated_history = [
+        *state["history"],
+        model_content.model_dump(exclude_none=True),
+    ]
+
+    pending_calls = [
+        {
+            "name": call.name,
+            "args": call.args or {},
+        }
+        for call in function_calls
+    ]
+
+    print(
+        "[Agent] "
+        + ", ".join(call["name"] for call in pending_calls)
+    )
+
+    return {
+        **state,
+        "history": updated_history,
+        "pending_calls": pending_calls,
+        "steps": state["steps"] + 1,
+    }
+
+
+async def execute_tools(state: AgentState) -> AgentState:
+    server_params = StdioServerParameters(
+        command=sys.executable,
+        args=["-m", "netdocai.mcp_server"],
+        env=os.environ.copy(),
+    )
+
+    updated_history = list(state["history"])
+    observations = list(state["observations"])
+
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write) as session:
+            await session.initialize()
+
+            for call in state["pending_calls"]:
+                name = call["name"]
+                arguments = call["args"]
+
+                print(f"[Tool] Executing {name}...")
+
+                result = await session.call_tool(
+                    name,
+                    arguments=arguments,
+                )
+
+                result_text = "\n".join(
+                    content.text
+                    for content in result.content
+                    if hasattr(content, "text") and content.text
+                )
+
+                if not result_text:
+                    result_text = "Tool returned no text result."
+
+                print(f"[Tool] {name} completed.")
+
+                observations.append(
+                    {
+                        "tool": name,
+                        "arguments": arguments,
+                        "result": result_text,
+                    }
+                )
+
+                function_response = types.Part.from_function_response(
+                    name=name,
+                    response={"result": result_text},
+                )
+
+                updated_history.append(
+                    types.Content(
+                        role="user",
+                        parts=[function_response],
+                    ).model_dump(exclude_none=True)
+                )
+
+    return {
+        **state,
+        "history": updated_history,
+        "observations": observations,
+        "pending_calls": [],
+    }
+
+
+async def assess(state: AgentState) -> AgentState:
+    observation_text = "\n\n".join(
+        f"Tool: {item['tool']}\n"
+        f"Arguments: {item['arguments']}\n"
+        f"Result:\n{item['result']}"
+        for item in state["observations"]
+    )
+
+    prompt = f"""
+You are assessing an ongoing network troubleshooting investigation.
+
+Problem:
+{state["problem"]}
+
+Observed evidence:
+{observation_text}
+
+Current hypothesis:
+{state["hypothesis"] or "None yet."}
+
+Based only on the available evidence:
+
+1. State the most likely current hypothesis.
+2. Identify the strongest supporting evidence.
+3. Identify what remains uncertain.
+4. Do not claim that a test was performed unless it appears in the evidence.
+
+Keep the assessment concise.
+"""
+
+    response = await generate_content(
+        contents=prompt,
+        config=types.GenerateContentConfig(),
+    )
+
+    return {
+        **state,
+        "hypothesis": response.text or "",
+    }
+
+
+async def finalize(state: AgentState) -> dict[str, Any]:
+    report = build_report(
+        problem=state["problem"],
+        observations=state["observations"],
+        hypothesis=state["hypothesis"],
+        conclusion=(
+            "The investigation reached the maximum number of diagnostic "
+            "steps before a final diagnosis was produced."
+        ),
+    )
+
+    return {
+        "response": report,
+    }
+
+
+def route_after_plan(state: AgentState) -> str:
+    if state["response"]:
+        return "done"
+
+    if state["steps"] > MAX_STEPS:
+        return "finalize"
+
+    return "execute"
+
+
+graph_builder = StateGraph(AgentState)
+
+graph_builder.add_node("plan", plan)
+graph_builder.add_node("execute_tools", execute_tools)
+graph_builder.add_node("assess", assess)
+graph_builder.add_node("finalize", finalize)
+
+graph_builder.add_edge(START, "plan")
+
+graph_builder.add_conditional_edges(
+    "plan",
+    route_after_plan,
+    {
+        "execute": "execute_tools",
+        "finalize": "finalize",
+        "done": END,
+    },
+)
+
+graph_builder.add_edge("execute_tools", "assess")
+graph_builder.add_edge("assess", "plan")
+graph_builder.add_edge("finalize", END)
+
+graph = graph_builder.compile()
+
+
+async def run_agent(problem: str) -> str:
+    tool_definitions = await get_tool_definitions()
+
+    initial_history = [
+        {
+            "role": "user",
+            "parts": [
+                {
+                    "text": (
+                        f"{SYSTEM_PROMPT}\n\n"
+                        f"{build_user_prompt(problem)}"
+                    )
+                }
+            ],
+        }
+    ]
+
+    result = await graph.ainvoke(
+        {
+            "problem": problem,
+            "history": initial_history,
+            "tool_definitions": tool_definitions,
+            "pending_calls": [],
+            "observations": [],
+            "hypothesis": "",
+            "response": "",
+            "steps": 0,
+        }
+    )
+
+    return result["response"]

netdocai/config.py

@@ -0,0 +1,44 @@
+import os
+from dataclasses import dataclass
+
+from dotenv import load_dotenv
+
+
+DEFAULT_GEMINI_MODEL = "gemini-3.6-flash"
+
+
+@dataclass(frozen=True)
+class Settings:
+    gemini_api_key: str
+    gemini_model: str = DEFAULT_GEMINI_MODEL
+
+    @classmethod
+    def from_env(cls) -> "Settings":
+        load_dotenv()
+
+        api_key = os.getenv("GEMINI_API_KEY", "").strip()
+
+        if not api_key:
+            raise RuntimeError(
+                "Missing required environment variable: GEMINI_API_KEY"
+            )
+
+        model = os.getenv(
+            "GEMINI_MODEL",
+            DEFAULT_GEMINI_MODEL,
+        )
+
+        return cls(
+            gemini_api_key=api_key,
+            gemini_model=model,
+        )
+
+
+class LazySettings:
+    def __getattr__(self, name: str):
+        return getattr(Settings.from_env(), name)
+
+
+settings = LazySettings()
+
+__all__ = ["Settings", "settings"]

netdocai/gemini.py

@@ -0,0 +1,44 @@
+import asyncio
+from typing import Any
+
+from google import genai
+from google.genai import types
+from google.genai.errors import ClientError, ServerError
+
+from .config import settings
+
+
+async def generate_content(
+    contents: Any,
+    config: types.GenerateContentConfig,
+    retries: int = 3,
+) -> Any:
+    async with genai.Client(api_key=settings.gemini_api_key).aio as client:
+        for attempt in range(retries):
+            try:
+                return await client.models.generate_content(
+                    model=settings.gemini_model,
+                    contents=contents,
+                    config=config,
+                )
+
+            except ServerError:
+                if attempt == retries - 1:
+                    raise
+
+            except ClientError as exc:
+                status_code = getattr(exc, "code", None)
+
+                if status_code != 429 or attempt == retries - 1:
+                    raise
+
+            delay = 2 ** attempt
+
+            print(
+                f"[Gemini] Temporary API failure. "
+                f"Retrying in {delay}s..."
+            )
+
+            await asyncio.sleep(delay)
+
+    raise RuntimeError("Gemini request failed after all retries.")

netdocai/knowledge/dns.md

@@ -0,0 +1,55 @@
+# DNS Troubleshooting
+
+DNS translates domain names into IP addresses.
+
+## Common Outcomes
+
+### Successful Resolution
+
+A hostname resolves when the DNS resolver returns one or more valid IP addresses.
+
+Check both IPv4 and IPv6 when troubleshooting connectivity:
+
+- A record -> IPv4 address
+- AAAA record -> IPv6 address
+
+### NXDOMAIN
+
+NXDOMAIN means the queried domain name does not exist according to the DNS server.
+
+Possible causes:
+- Typo in hostname
+- Missing DNS record
+- Incorrect domain configuration
+
+### SERVFAIL
+
+SERVFAIL means the DNS server failed to provide a valid answer.
+
+Possible causes:
+- Upstream DNS failure
+- DNSSEC validation failure
+- Authoritative nameserver problems
+- Temporary resolver problems
+
+### DNS Timeouts
+
+A timeout means the DNS query did not receive a response within the expected period.
+
+Possible causes:
+- Network connectivity problems
+- Unreachable resolver
+- Firewall filtering
+- Resolver overload
+
+### Troubleshooting
+
+When investigating DNS problems:
+
+1. Check whether the hostname resolves.
+2. Check both A and AAAA records.
+3. Measure DNS response time.
+4. Compare results using another resolver.
+5. Check whether the problem is intermittent or persistent.
+
+DNS problems can appear as application failures even when TCP and HTTP connectivity are otherwise healthy.

netdocai/knowledge/http.md

@@ -0,0 +1,46 @@
+# HTTP Troubleshooting
+
+HTTP operates above the transport layer and is used by web applications and APIs.
+
+## HTTP Status Codes
+
+### 2xx
+
+Successful request.
+
+### 4xx
+
+Client-side request error.
+
+### 5xx
+
+Server-side error.
+
+## Response Time
+
+HTTP response time measures how long the client waits for the server response.
+
+High response time can result from:
+
+- Slow application processing
+- Slow upstream services
+- Server overload
+- Network latency
+- Backend database delays
+
+## Time To First Byte
+
+TTFB measures the time until the first byte of the response is received.
+
+High TTFB can indicate that the server or an upstream dependency is taking a long time to begin generating the response.
+
+## Troubleshooting
+
+When investigating HTTP slowness:
+
+1. Check the HTTP status code.
+2. Measure total response time.
+3. Measure TTFB when available.
+4. Compare multiple requests.
+5. Compare HTTP and HTTPS behavior.
+6. Determine whether the issue is persistent or intermittent.

netdocai/knowledge/linux_networking.md

@@ -0,0 +1,51 @@
+# Linux Networking Troubleshooting
+
+Linux provides several tools for investigating network connectivity.
+
+## Ping
+
+Ping measures ICMP reachability and round-trip latency.
+
+Useful metrics include:
+
+- Packet loss
+- Minimum latency
+- Average latency
+- Maximum latency
+- Jitter
+
+Ping does not prove that TCP or HTTP services are working.
+
+## DNS
+
+Common commands include:
+
+- getent hosts
+- dig
+- nslookup
+
+Use DNS tools to determine whether hostname resolution is working correctly.
+
+## TCP
+
+Useful tools include:
+
+- nc
+- curl
+- ss
+
+A TCP connection test checks whether a destination port can be reached.
+
+## HTTP
+
+curl can be used to test HTTP and HTTPS endpoints and inspect:
+
+- Status code
+- Response headers
+- Response timing
+
+## Traceroute
+
+Traceroute can help identify routing paths and where latency or packet loss may occur.
+
+A single failed or missing traceroute hop does not necessarily mean that the network path is broken because routers may filter or deprioritize diagnostic traffic.

netdocai/knowledge/tcp.md

@@ -0,0 +1,42 @@
+# TCP Troubleshooting
+
+TCP provides reliable, connection-oriented communication between applications.
+
+## TCP Connection
+
+A successful TCP connection indicates that a client was able to establish a connection to the destination host and port.
+
+For HTTPS, the usual destination port is 443.
+
+## Connection Failures
+
+Common causes include:
+
+- Destination service is not listening
+- Firewall filtering
+- Network routing problems
+- Server unavailable
+- Incorrect port
+
+## Connection Latency
+
+Measure the time required to establish a TCP connection.
+
+High TCP connection latency can indicate:
+
+- Network congestion
+- Routing problems
+- High server load
+- Packet loss and retransmissions
+
+## Troubleshooting
+
+When investigating TCP problems:
+
+1. Confirm the destination hostname resolves.
+2. Attempt a TCP connection to the expected port.
+3. Measure connection latency.
+4. Compare results across repeated attempts.
+5. Check whether failures are intermittent or persistent.
+
+A successful TCP connection does not prove that the application itself is healthy.

netdocai/main.py

@@ -0,0 +1,85 @@
+import asyncio
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Prompt
+
+from .onboarding import get_api_key
+
+
+console = Console()
+
+
+def main() -> None:
+    console.print(
+        Panel(
+            "[bold]NetDocAI[/bold]\n"
+            "Autonomous Network Troubleshooter",
+            expand=False,
+        )
+    )
+
+    console.print()
+
+    try:
+        get_api_key()
+    except (KeyboardInterrupt, EOFError):
+        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        return
+    except OSError as exc:
+        console.print(f"\n[red]Configuration error: {exc}[/red]")
+        return
+
+    from .agent import run_agent
+
+    console.print()
+
+    problem = Prompt.ask(
+        "[bold]Describe your network problem[/bold]"
+    ).strip()
+
+    if not problem:
+        console.print("[red]No problem provided.[/red]")
+        return
+
+    console.print()
+    console.print(
+        Panel(
+            problem,
+            title="Problem",
+            expand=False,
+        )
+    )
+
+    console.print()
+    console.print("[bold cyan]Investigating...[/bold cyan]")
+    console.print()
+
+    try:
+        with console.status(
+            "[bold cyan]NetDocAI is investigating...[/bold cyan]",
+            spinner="dots",
+        ):
+            response = asyncio.run(run_agent(problem))
+
+        console.print()
+        console.print(
+            Panel(
+                response,
+                title="NetDocAI Diagnostic Report",
+                border_style="green",
+            )
+        )
+    except Exception as exc:
+        console.print()
+        console.print(
+            Panel(
+                f"[red]{exc}[/red]",
+                title="Agent Error",
+                border_style="red",
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()