puli-plg 0.1.26__py3-none-any.whl

puli_mcp_server/mcp_server/server.py

@@ -0,0 +1,123 @@
+import argparse
+import asyncio
+import json
+import logging
+from typing import List
+from mcp.server.fastmcp import FastMCP, Context
+
+from puli_mcp_server.mcp_server.models import ChangeSet
+from puli_mcp_server.proxy_client import ProxyClient
+from puli_mcp_server.embedding_client.client import EmbeddingClient
+from puli_mcp_server.embedding_client.config import EmbeddingConfig
+from puli_mcp_server.llm_agent.llm_agent import LLMAgent
+from puli_mcp_server.llm_agent.config import LLMAgentConfig
+from puli_mcp_server.llm_agent.models import LLMQueryRequest
+
+logger = logging.getLogger(__name__)
+
+mcp = FastMCP("code-reviewer")
+
+# Initialize clients with remote configuration
+proxy_client = ProxyClient()
+
+# Initialize embedding client with remote config
+embedding_config = EmbeddingConfig.from_remote(proxy_client.mcp_config["config"])
+embedding_client = EmbeddingClient(config=embedding_config)
+
+# Initialize LLM agent with remote config and prompts
+llm_config = LLMAgentConfig.from_remote(
+    proxy_client.mcp_config["config"],
+    proxy_client.mcp_config["prompts"]
+)
+llm_agent = LLMAgent(config=llm_config)
+
+
+def log_configuration():
+    """Log MCP configuration after logging is properly initialized."""
+    logger.info("=" * 60)
+    logger.info("MCP Configuration pulled from proxy:")
+    logger.info("  Config keys: %s", list(proxy_client.mcp_config.get("config", {}).keys()))
+    logger.info("  Prompt keys: %s", list(proxy_client.mcp_config.get("prompts", {}).keys()))
+    logger.info("=" * 60)
+    logger.info("Embedding Config - Model: %s", embedding_config.model)
+    logger.info("LLM Config - Provider: %s, Model: %s, Temperature: %s", 
+                llm_config.provider, llm_config.model, llm_config.temperature)
+    logger.info("System prompt length: %d characters", len(llm_config.system_prompt))
+    logger.info("=" * 60)
+
+
+# --- 3. Define the Tool ---
+
+@mcp.tool()
+async def puli_herd(change_sets: List[ChangeSet], ctx: Context) -> str:
+    """
+    Puli Guard is a tool that analyzes a list of code changes for security risks, complexity, and
+    infrastructure impact by comparing them with historical incidents.
+    """
+
+    results = []
+
+    for i, change_set in enumerate(change_sets, 1):
+        await ctx.info(f"Analyzing change set {i}: {change_set.goal}")
+
+        # 1. Create embedding vector
+        embedding_text = change_set.to_embedding_string()
+        vector = embedding_client.generate_embedding(embedding_text)
+
+        # 2. Query proxy for similar historical incidents and relevant chaos patterns (in parallel)
+        similar_incidents, similar_chaos_patterns = await asyncio.gather(
+            asyncio.to_thread(proxy_client.search_incidents, query_vector=vector),
+            asyncio.to_thread(proxy_client.search_chaos_patterns, query_vector=vector)
+        )
+
+        # 3. Create LLM query request and get analysis
+        query_request = LLMQueryRequest(
+            change_set=change_set,
+            historical_incidents=similar_incidents,
+            relevant_chaos_patterns=similar_chaos_patterns
+        )
+        risk_assessment = await llm_agent.query(query_request)
+        results.append(risk_assessment.to_str())
+
+    return "\n\n---\n\n".join(results)
+
+
+class MockContext:
+    """Mock context for CLI mode that prints info messages to stdout."""
+    async def info(self, msg: str):
+        print(f"[INFO] {msg}")
+
+
+async def run_from_file(file_path: str) -> str:
+    """Run analysis from a JSON file containing change sets."""
+    with open(file_path) as f:
+        data = json.load(f)
+
+    change_sets = [ChangeSet(**cs) for cs in data]
+    return await puli_herd(change_sets, MockContext())
+
+
+def main():
+    # Configure logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='[%(asctime)s] %(levelname)s - %(name)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    
+    # Log configuration after logging is set up
+    log_configuration()
+    
+    parser = argparse.ArgumentParser(description="Puli Code Reviewer")
+    parser.add_argument("--file", "-f", help="JSON file with change sets (CLI mode)")
+    args = parser.parse_args()
+
+    if args.file:
+        result = asyncio.run(run_from_file(args.file))
+        print(result)
+    else:
+        mcp.run()  # Original MCP mode
+
+
+if __name__ == "__main__":
+    main()

puli_mcp_server/proxy_client/__init__.py

@@ -0,0 +1,3 @@
+from .client import ProxyClient
+
+__all__ = ["ProxyClient"]

puli_mcp_server/proxy_client/client.py

@@ -0,0 +1,264 @@
+from typing import List, Optional, Dict, Any
+import urllib.request
+import urllib.error
+import json
+import logging
+
+from puli_models import IncidentQueryResult, ChaosPatternQueryResult
+
+from .config import ProxyConfig
+from .token_manager import get_machine_id, read_token, write_token, clear_token
+
+logger = logging.getLogger(__name__)
+
+
+class ProxyClient:
+    """Client for interacting with the Puli proxy service."""
+
+    def __init__(self, config: Optional[ProxyConfig] = None):
+        """
+        Initialize the proxy client.
+
+        Args:
+            config: ProxyConfig instance. If None, loads from environment variables.
+        """
+        self.config = config or ProxyConfig.from_env()
+        self._ensure_authenticated()
+        self.mcp_config = self._fetch_mcp_config()
+
+    def _ensure_authenticated(self) -> None:
+        """Load existing JWT or register to obtain one."""
+        token = read_token()
+        if token:
+            self.config.jwt_token = token
+            return
+        self._register()
+
+    def _register(self) -> None:
+        """Register this machine and store the returned JWT."""
+        machine_id = get_machine_id()
+        url = f"{self.config.base_url}/auth/register"
+        headers = {
+            "X-API-Key": self.config.api_key,
+            "Content-Type": "application/json",
+        }
+        if self.config.identity_token:
+            headers["Authorization"] = f"Bearer {self.config.identity_token}"
+
+        body = json.dumps({"machine_id": machine_id}).encode()
+        req = urllib.request.Request(url, data=body, headers=headers, method="POST")
+
+        try:
+            with urllib.request.urlopen(req) as response:
+                result = json.loads(response.read().decode())
+        except urllib.error.HTTPError as e:
+            error_body = e.read().decode() if e.fp else ""
+            raise RuntimeError(f"Registration failed: {e.code} - {error_body}") from e
+
+        token = result["token"]
+        write_token(token)
+        self.config.jwt_token = token
+
+    def _fetch_mcp_config(self) -> Dict[str, Any]:
+        """Fetch MCP configuration from the proxy after authentication."""
+        try:
+            return self._request("GET", "/config/mcp")
+        except RuntimeError as e:
+            # If config fetch fails, return empty config (allows local dev fallback)
+            if "404" in str(e) or "500" in str(e):
+                logger.warning("Failed to fetch MCP config from proxy: %s", e)
+                logger.warning("Falling back to empty config. Client will use from_env() methods for local development.")
+                return {"config": {}, "prompts": {}}
+            raise
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        data: Optional[dict] = None,
+    ) -> dict:
+        """Make an HTTP request to the proxy with auto-retry on 401."""
+        url = f"{self.config.base_url}{path}"
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        if self.config.jwt_token:
+            headers["X-Auth-Token"] = self.config.jwt_token
+
+        if self.config.identity_token:
+            headers["Authorization"] = f"Bearer {self.config.identity_token}"
+
+        body = json.dumps(data).encode() if data else None
+        request = urllib.request.Request(url, data=body, headers=headers, method=method)
+
+        try:
+            with urllib.request.urlopen(request) as response:
+                return json.loads(response.read().decode())
+        except urllib.error.HTTPError as e:
+            if e.code == 401:
+                # Token expired or invalid — re-register once and retry
+                clear_token()
+                self._register()
+                return self._retry_request(method, url, data)
+            error_body = e.read().decode() if e.fp else ""
+            raise RuntimeError(f"Proxy request failed: {e.code} - {error_body}") from e
+
+    def _retry_request(
+        self,
+        method: str,
+        url: str,
+        data: Optional[dict] = None,
+    ) -> dict:
+        """Single retry after re-registration (no infinite loop)."""
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        if self.config.jwt_token:
+            headers["X-Auth-Token"] = self.config.jwt_token
+
+        if self.config.identity_token:
+            headers["Authorization"] = f"Bearer {self.config.identity_token}"
+
+        body = json.dumps(data).encode() if data else None
+        request = urllib.request.Request(url, data=body, headers=headers, method=method)
+
+        try:
+            with urllib.request.urlopen(request) as response:
+                return json.loads(response.read().decode())
+        except urllib.error.HTTPError as e:
+            error_body = e.read().decode() if e.fp else ""
+            raise RuntimeError(f"Proxy request failed after retry: {e.code} - {error_body}") from e
+
+    def _search(
+        self,
+        query_vector: List[float],
+        url: str,
+        limit: int = 10,
+        filter_expr: Optional[str] = None,
+        output_fields: Optional[List[str]] = None,
+    ) -> List[Any]:
+        """
+        Search for similar incidents by vector.
+
+        Args:
+            query_vector: The embedding vector to search with.
+            limit: Maximum number of results to return.
+            filter_expr: Optional filter expression.
+            output_fields: Fields to include in results.
+
+        Returns:
+            List of matching IncidentQueryResult objects.
+        """
+        data = {
+            "query_vector": query_vector,
+            "limit": limit,
+        }
+        if filter_expr:
+            data["filter"] = filter_expr
+        if output_fields:
+            data["output_fields"] = output_fields
+
+        response = self._request("POST", url, data)
+        return response.get("results", [])
+
+    def search_incidents(
+        self,
+        query_vector: List[float],
+        limit: int = 10,
+        filter_expr: Optional[str] = None,
+        output_fields: Optional[List[str]] = None,
+    ) -> List[IncidentQueryResult]:
+        results = self._search(query_vector, "/incidents/search", limit, filter_expr, output_fields)
+        return [IncidentQueryResult.from_dict(r) for r in results]
+
+    def search_chaos_patterns(
+        self,
+        query_vector: List[float],
+        limit: int = 10,
+        filter_expr: Optional[str] = None,
+        output_fields: Optional[List[str]] = None,
+    ) -> List[ChaosPatternQueryResult]:
+        results = self._search(query_vector, "/chaos-patterns/search", limit, filter_expr, output_fields)
+        return [ChaosPatternQueryResult.from_dict(r) for r in results]
+
+    def query(
+        self,
+        filter_expr: str,
+        output_fields: Optional[List[str]] = None,
+        limit: int = 100,
+    ) -> List[IncidentQueryResult]:
+        """
+        Query incidents by filter expression (no vector search).
+
+        Args:
+            filter_expr: Filter expression.
+            output_fields: Fields to include in results.
+            limit: Maximum number of results.
+
+        Returns:
+            List of matching IncidentQueryResult objects.
+        """
+        data = {
+            "filter": filter_expr,
+            "limit": limit,
+        }
+        if output_fields:
+            data["output_fields"] = output_fields
+
+        response = self._request("POST", "/incidents/query", data)
+        return [IncidentQueryResult.from_dict(r) for r in response.get("results", [])]
+
+    def insert(self, incidents: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Insert incidents into the collection.
+
+        Args:
+            incidents: List of incident dictionaries to insert.
+
+        Returns:
+            Insert result with IDs of inserted records.
+        """
+        data = {"data": incidents}
+        return self._request("POST", "/incidents/insert", data)
+
+    def get_by_id(self, incident_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get a single incident by ID.
+
+        Args:
+            incident_id: The incident ID to retrieve.
+
+        Returns:
+            Incident dictionary or None if not found.
+        """
+        try:
+            return self._request("GET", f"/incidents/{incident_id}")
+        except RuntimeError as e:
+            if "404" in str(e):
+                return None
+            raise
+
+    def delete(self, ids: List[str]) -> Dict[str, Any]:
+        """
+        Delete incidents by IDs.
+
+        Args:
+            ids: List of incident IDs to delete.
+
+        Returns:
+            Delete result.
+        """
+        data = {"ids": ids}
+        return self._request("DELETE", "/incidents", data)
+
+    def count(self) -> int:
+        """
+        Get the total number of incidents in the collection.
+
+        Returns:
+            Number of incidents.
+        """
+        response = self._request("GET", "/incidents/count")
+        return response.get("count", 0)

puli_mcp_server/proxy_client/config.py

@@ -0,0 +1,74 @@
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from google.auth.transport.requests import Request
+from google.oauth2 import service_account
+
+# Bundled service account key path (inside the package)
+_CREDENTIALS_PATH = Path(__file__).parent.parent / "credentials" / "service-account.json"
+
+# Default Cloud Run URL
+_DEFAULT_PROXY_URL = "https://puli-proxy-gpvaoh5hka-uw.a.run.app"
+
+
+def _load_credentials(target_audience: str) -> service_account.IDTokenCredentials | None:
+    """Load GCP credentials from the bundled service account key."""
+    sa_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", str(_CREDENTIALS_PATH))
+
+    if not Path(sa_path).exists():
+        return None
+
+    return service_account.IDTokenCredentials.from_service_account_file(
+        sa_path, target_audience=target_audience
+    )
+
+
+@dataclass
+class ProxyConfig:
+    """Configuration for the proxy client."""
+
+    base_url: str
+    api_key: str
+    _credentials: service_account.IDTokenCredentials | None = field(
+        default=None, repr=False
+    )
+    _jwt_token: str | None = field(default=None, repr=False)
+
+    @property
+    def identity_token(self) -> str | None:
+        """Get a valid identity token, refreshing automatically if expired."""
+        if self._credentials is None:
+            return None
+
+        if not self._credentials.valid:
+            self._credentials.refresh(Request())
+
+        return self._credentials.token
+
+    @property
+    def jwt_token(self) -> str | None:
+        return self._jwt_token
+
+    @jwt_token.setter
+    def jwt_token(self, value: str | None) -> None:
+        self._jwt_token = value
+
+    @classmethod
+    def from_env(cls) -> "ProxyConfig":
+        """Load configuration from environment variables."""
+        base_url = os.environ.get("PROXY_BASE_URL", _DEFAULT_PROXY_URL)
+        if not base_url:
+            raise ValueError("PROXY_BASE_URL environment variable is required")
+
+        api_key = os.environ.get("PROXY_API_KEY", "puli-proxy-api-key")
+        if not api_key:
+            raise ValueError("PROXY_API_KEY environment variable is required")
+
+        credentials = _load_credentials(target_audience=base_url)
+
+        return cls(
+            base_url=base_url.rstrip("/"),
+            api_key=api_key,
+            _credentials=credentials,
+        )

puli_mcp_server/proxy_client/token_manager.py

@@ -0,0 +1,36 @@
+import uuid
+from pathlib import Path
+
+_PULI_DIR = Path.home() / ".puli"
+_MACHINE_ID_FILE = _PULI_DIR / "machine_id"
+_TOKEN_FILE = _PULI_DIR / "token"
+
+
+def _ensure_dir() -> None:
+    _PULI_DIR.mkdir(mode=0o700, exist_ok=True)
+
+
+def get_machine_id() -> str:
+    _ensure_dir()
+    if _MACHINE_ID_FILE.exists():
+        return _MACHINE_ID_FILE.read_text().strip()
+    machine_id = str(uuid.uuid4())
+    _MACHINE_ID_FILE.write_text(machine_id)
+    return machine_id
+
+
+def read_token() -> str | None:
+    if _TOKEN_FILE.exists():
+        token = _TOKEN_FILE.read_text().strip()
+        return token or None
+    return None
+
+
+def write_token(token: str) -> None:
+    _ensure_dir()
+    _TOKEN_FILE.write_text(token)
+
+
+def clear_token() -> None:
+    if _TOKEN_FILE.exists():
+        _TOKEN_FILE.unlink()

puli_mcp_server/test_diff.json

@@ -0,0 +1,16 @@
+{
+    "change_sets": [
+      {
+        "goal": "Add chaos_scenario field to RiskAssessment model to track chaos testing scenarios",
+        "changes": [
+          {
+            "file_path": "src/puli_mcp_server/llm_agent/models.py",
+            "change_type": "modify",
+            "diff_content": "@@ -87,6 +87,11 @@ class RiskAssessment(BaseModel):\n     technical_finding: TechnicalFinding\n     business_context: BusinessContext\n     \n     consequence: str = Field(\n         ..., \n         description=\"Description: What happens to the user or the business. Style Rule: Be strictly factual. No drama. No hyperbole. Example: \\\"User is double-charged. Support ticket generated.\\\" (NOT \\\"Catastrophic failure destroys trust\\\").\"\n     )\n+\n+    chaos_scenario: Optional[str] = Field(\n+        None,\n+        description=\"Description of the chaos scenario that was run to test this risk assessment.\"\n+    )\n     \n     historical_incident: Optional[RealIncident] = Field(\n         None, \n         description=\"Only populate if a famous/known incident matches this exact failure pattern.\"\n     )"
+          }
+        ],
+        "related_infrastructure": "Pydantic Models, LLM Agent",
+        "additional_context": "Added optional field to track chaos engineering scenarios used in risk assessment validation"
+      }
+    ]
+  }

puli_models/__init__.py

@@ -0,0 +1,9 @@
+from .incidents import Incident, IncidentQueryResult
+from .chaos_patterns import ChaosPattern, ChaosPatternQueryResult
+
+__all__ = [
+    "Incident",
+    "IncidentQueryResult",
+    "ChaosPattern",
+    "ChaosPatternQueryResult",
+]

puli_models/chaos_patterns.py

@@ -0,0 +1,89 @@
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class ChaosPattern:
+    """Represents a chaos pattern record in the Zilliz collection."""
+
+    id: str
+    vector: List[float]
+    name: str
+    category: str
+    description: str
+    cause: str
+    symptoms: str
+    technology: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ChaosPattern":
+        """Create a ChaosPattern from a dictionary."""
+        return cls(
+            id=data.get("id"),
+            vector=list(data.get("vector", [])),
+            name=data.get("name"),
+            category=data.get("category"),
+            description=data.get("description"),
+            cause=data.get("cause"),
+            symptoms=data.get("symptoms"),
+            technology=data.get("technology"),
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for insertion."""
+        return {
+            "id": self.id,
+            "vector": self.vector,
+            "name": self.name,
+            "category": self.category,
+            "description": self.description,
+            "cause": self.cause,
+            "symptoms": self.symptoms,
+            "technology": self.technology,
+        }
+
+
+@dataclass
+class ChaosPatternQueryResult:
+    """Represents a lightweight chaos pattern record for search/query results."""
+
+    name: str
+    category: str
+    description: str
+    cause: str
+    symptoms: str
+    technology: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ChaosPatternQueryResult":
+        """Create a ChaosPatternQueryResult from a dictionary."""
+        return cls(
+            name=data.get("name"),
+            category=data.get("category"),
+            description=data.get("description"),
+            cause=data.get("cause"),
+            symptoms=data.get("symptoms"),
+            technology=data.get("technology"),
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for serialization."""
+        return {
+            "name": self.name,
+            "category": self.category,
+            "description": self.description,
+            "cause": self.cause,
+            "symptoms": self.symptoms,
+            "technology": self.technology,
+        }
+
+    def to_prompt_str(self) -> str:
+        """Returns a string representation of the chaos pattern query result."""
+        prompt = f"Title: {self.name}\n" \
+            + f"Category: {self.category}\n" \
+            + f"Description: {self.description}\n" \
+            + f"Root Causes: {self.cause}\n" \
+            + f"Common Symptoms: {self.symptoms}\n"
+        if self.technology:
+            prompt += f"Technology: {self.technology}\n"
+        return prompt