naive-knowledge-base 0.1.0__py3-none-any.whl

api_models/flow_api_model.py

@@ -0,0 +1,390 @@
+from copy import deepcopy
+from typing import Dict, List, Optional, Union, Any, ClassVar, Type
+import requests
+from smolagents import ApiModel, ChatMessage, MessageRole
+from smolagents.tools import Tool
+from smolagents.models import get_clean_message_list, tool_role_conversions
+from openai.types.chat import ChatCompletion
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+def _request_access_token() -> str:
+    """Request an access token from the Flow API"""
+    response = requests.post(
+        "https://flow.ciandt.com/auth-engine-api/v1/api-key/token",
+        headers={"Content-Type": "application/json", "FlowTenant": os.getenv("FLOW_TENANT", "flowteam")},
+        json={"clientId": os.getenv("FLOW_CLIENT_ID"), "clientSecret": os.getenv("FLOW_CLIENT_SECRET"), "appToAccess": "llm-api"}
+    )
+
+    return response.json().get("access_token")
+
+_access_token = _request_access_token()
+
+class ModelConfig:
+    """Configuration for a specific model type"""
+    def __init__(
+        self,
+        api_url: str,
+        response_handler: str,
+        default_params: Dict[str, Any],
+        requires_anthropic_format: bool = False,
+        requires_system_message_separation: bool = False,
+    ):
+        self.api_url = api_url
+        self.response_handler = response_handler  # 'claude' or 'openai'
+        self.default_params = default_params
+        self.requires_anthropic_format = requires_anthropic_format
+        self.requires_system_message_separation = requires_system_message_separation
+
+
+class FlowApiModel(ApiModel):
+    """
+    A unified API model wrapper for Flow-based models (Claude, GPT, DeepSeek),
+    providing a simple interface to interact with different model types.
+    """
+    # Registry of model-specific configurations
+    MODEL_REGISTRY: ClassVar[Dict[str, ModelConfig]] = {
+        # Claude models
+        "anthropic.claude-37-sonnet": ModelConfig(
+            api_url="https://flow.ciandt.com/ai-orchestration-api/v1/bedrock/invoke",
+            response_handler="claude",
+            default_params={
+                "anthropic_version": "bedrock-2023-05-31",
+                "max_tokens": 20000,
+                "temperature": 0.5,
+                "top_p": 0.9,
+                "top_k": 250,
+            },
+            requires_anthropic_format=True,
+            requires_system_message_separation=True,
+        ),
+        
+        # DeepSeek models
+        "DeepSeek-R1": ModelConfig(
+            api_url="https://flow.ciandt.com/ai-orchestration-api/v1/foundry/chat/completions",
+            response_handler="openai",
+            default_params={
+                "max_tokens": 8192,
+            },
+        ),
+        
+        # GPT models
+        "o3-mini": ModelConfig(
+            api_url="https://flow.ciandt.com/ai-orchestration-api/v1/openai/chat/completions",
+            response_handler="openai",
+            default_params={
+                "max_tokens": 8192,
+                "top_p": 1.0,
+                "reasoning_effort": "medium",
+                "presence_penalty": 0,
+                "frequency_penalty": 0,
+            },
+        ),
+
+        "gpt-4o": ModelConfig(
+            api_url="https://flow.ciandt.com/ai-orchestration-api/v1/openai/chat/completions",
+            response_handler="openai",
+            default_params={
+                "max_tokens": 8192,
+                "temperature": 0.7,
+                "top_p": 1.0,
+                "presence_penalty": 0,
+                "frequency_penalty": 0,
+            },
+        ),
+
+        "gpt-4.1": ModelConfig(
+            api_url="https://flow.ciandt.com/ai-orchestration-api/v1/openai/chat/completions",
+            response_handler="openai",
+            default_params={
+                "max_tokens": 8192,
+                "temperature": 0.7,
+                "top_p": 1.0,
+                "presence_penalty": 0,
+                "frequency_penalty": 0,
+            },
+        ),
+
+        "o1-mini": ModelConfig(
+            api_url="https://flow.ciandt.com/ai-orchestration-api/v1/openai/chat/completions",
+            response_handler="openai",
+            default_params={
+                "reasoning_effort": "medium",
+                "top_p": 1.0,
+                "presence_penalty": 0,
+                "frequency_penalty": 0,
+            },
+        ),
+    }
+
+    @classmethod
+    def get_model_config(cls, model_id: str) -> ModelConfig:
+        """Get configuration for a specific model ID"""
+        if model_id not in cls.MODEL_REGISTRY:
+            raise ValueError(f"Unsupported model: {model_id}. Available models: {list(cls.MODEL_REGISTRY.keys())}")
+        return cls.MODEL_REGISTRY[model_id]
+
+    @classmethod
+    def create(cls, model_id: str, **kwargs) -> 'FlowApiModel':
+        """Factory method to create the appropriate model instance"""
+        return cls(model_id=model_id, **kwargs)
+
+    def __init__(
+        self,
+        model_id: str,
+        api_key: str = _access_token,
+        flow_tenant: str = os.getenv("FLOW_TENANT", "flowteam"),
+        flow_agent: str = "feature_refinement_agent",
+        flow_operation_id: str = "default_operation",
+        **kwargs,
+    ):
+        super().__init__(model_id=model_id, **kwargs)
+        
+        # Get model-specific configuration
+        self.config = self.get_model_config(model_id)
+        
+        # Common attributes for all models
+        self.api_key = api_key
+        self.flow_tenant = flow_tenant
+        self.flow_agent = flow_agent
+        self.flow_operation_id = flow_operation_id
+        self.model_name = model_id  # For backward compatibility
+        
+        # Load model-specific default parameters
+        for key, value in self.config.default_params.items():
+            setattr(self, key, kwargs.get(key, value))
+        
+        # Store additional kwargs
+        self.kwargs = kwargs
+        
+        # Use model-specific API URL
+        self.api_url = kwargs.get("api_url", self.config.api_url)
+
+    def create_client(self):
+        """No additional client setup is required for Flow models."""
+        return None
+
+    def _create_headers(self) -> dict:
+        """
+        Creates the headers for the API request.
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "FlowTenant": self.flow_tenant,
+            "FlowAgent": self.flow_agent,
+            "Authorization": f"Bearer {self.api_key}",
+        }
+        
+        # Only add FlowOperationId if it's required (for Claude and GPT models)
+        if hasattr(self, "flow_operation_id"):
+            headers["FlowOperationId"] = self.flow_operation_id
+            
+        return headers
+
+    def _claude_get_tool_json_schema(self, tool: Tool) -> dict:
+        """Convert a Tool object to the JSON schema format expected by the API"""
+        properties = deepcopy(tool.inputs)
+        required = []
+        for key, value in properties.items():
+            if value["type"] == "any":
+                value["type"] = "string"
+            if not ("nullable" in value and value["nullable"]):
+                required.append(key)
+        return {
+            "name": tool.name,
+            "type": "function",
+            "description": tool.description,
+            "input_schema": {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+            },
+        }
+
+    def _prepare_payload(
+        self,
+        messages: List[Dict[str, str]],
+        stop_sequences: Optional[List[str]] = None,
+        grammar: Optional[str] = None,
+        tools_to_call_from: Optional[List[Tool]] = None,
+        custom_role_conversions: Optional[Dict[str, str]] = None,
+        convert_images_to_image_urls: bool = False,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Prepare the payload for the API request based on model type"""
+        # Handle system message separately for models that require it
+        system_message = None
+        if self.config.requires_system_message_separation:
+            system_message = next((message for message in messages if message["role"] == "system"), None)
+            messages = [message for message in messages if message["role"] != "system"]
+
+        # Prepare completion kwargs with base model params
+        completion_kwargs = self._prepare_completion_kwargs(
+            messages=messages,
+            stop_sequences=stop_sequences,
+            grammar=grammar,
+            tools_to_call_from=tools_to_call_from,
+            custom_role_conversions=custom_role_conversions,
+            convert_images_to_image_urls=convert_images_to_image_urls,
+            **kwargs,
+        )
+
+        # Add model-specific parameters
+        model_params = {}
+        for key, default_value in self.config.default_params.items():
+            model_params[key] = getattr(self, key, default_value)
+            
+        # Update with model ID and allowed models
+        model_params.update({
+            "allowedModels": [self.model_id],
+        })
+
+        if not self.config.response_handler == "claude":
+            model_params.update({
+                "model": self.model_id,
+            })
+        
+        # Special handling for o3-mini model
+        if self.model_id == "o3-mini" or self.model_id == "o1-mini":
+            if "max_tokens" in model_params:
+                del model_params["max_tokens"]
+        
+        # Update the completion kwargs with model params
+        completion_kwargs.update(model_params)
+
+        # del completion_kwargs["grammar"]  # Remove grammar from payload if it exists
+        
+        # Remove tool_choice from payload if it exists
+        if "tool_choice" in completion_kwargs:
+            del completion_kwargs["tool_choice"]
+            
+        # Special handling for Claude models - rename stop_sequences to stop
+        if self.config.response_handler == "claude" and "stop" in completion_kwargs:
+            completion_kwargs.pop("stop", None)
+
+        # Add tools if provided
+        if tools_to_call_from and self.config.requires_anthropic_format:
+            completion_kwargs["tools"] = [self._claude_get_tool_json_schema(tool) for tool in tools_to_call_from]
+
+        # Add system message if separated
+        if system_message and self.config.requires_system_message_separation:
+            completion_kwargs["system"] = system_message["content"]
+
+        return completion_kwargs
+
+    def _handle_response(self, response: requests.Response) -> ChatMessage:
+        """Handle the API response based on model type"""
+        if response.status_code != 200:
+            raise ValueError(
+                f"API call failed with status code {response.status_code}: {response.text}"
+            )
+        
+        response_data = response.json()
+        
+        # Update token counts
+        usage = response_data.get("usage") or {}
+        self._last_input_token_count = usage.get("input_tokens", 0)
+        self._last_output_token_count = usage.get("output_tokens", 0)
+        
+        if self.config.response_handler == "claude":
+            # Claude-style response handling
+            content_list = response_data.get("content", [])
+            if not content_list:
+                raise ValueError("API response contains no content.")
+
+            first_message_content = next(
+                (item["text"] for item in content_list if item.get("type") == "text"), None
+            )
+
+            if first_message_content is None:
+                raise ValueError("No valid text content found in the API response.")
+
+            return ChatMessage(
+                role=MessageRole.ASSISTANT,
+                content=first_message_content,
+                raw=response_data,
+            )
+        else:
+            # OpenAI-style response handling
+            chatCompletion = ChatCompletion.model_validate(response_data)
+            
+            return ChatMessage.from_dict(
+                chatCompletion.choices[0].message.model_dump(include={"role", "content", "tool_calls"}),
+                raw=response_data,
+            )
+
+    def generate(
+        self,
+        messages: List[Dict[str, Union[str, List[Dict]]]],
+        stop_sequences: Optional[List[str]] = None,
+        grammar: Optional[str] = None,
+        tools_to_call_from: Optional[List[Tool]] = None,
+        custom_role_conversions: Optional[Dict[str, str]] = None,
+        convert_images_to_image_urls: bool = False,
+        **kwargs,
+    ) -> ChatMessage:
+        """Process the input messages and return the model's response.
+
+        Parameters:
+            messages (`list[dict[str, str]]`):
+                A list of message dictionaries to be processed. Each dictionary should have the structure `{"role": "user/system", "content": "message content"}`.
+            stop_sequences (`List[str]`, *optional*):
+                A list of strings that will stop the generation if encountered in the model's output.
+            grammar (`str`, *optional*):
+                The grammar or formatting structure to use in the model's response.
+            tools_to_call_from (`List[Tool]`, *optional*):
+                A list of tools that the model can use to generate responses.
+            custom_role_conversions (`Dict[str, str]`, *optional*):
+                A dictionary mapping custom role names to standard role names.
+            convert_images_to_image_urls (`bool`, *optional*):
+                Whether to convert images in messages to URLs.
+            **kwargs:
+                Additional keyword arguments to be passed to the underlying model.
+
+        Returns:
+            `ChatMessage`: A chat message object containing the model's response.
+        """
+        # Prepare the payload
+        payload = self._prepare_payload(
+            messages=messages,
+            stop_sequences=stop_sequences,
+            grammar=grammar,
+            tools_to_call_from=tools_to_call_from,
+            custom_role_conversions=custom_role_conversions,
+            convert_images_to_image_urls=convert_images_to_image_urls,
+            **kwargs,
+        )
+
+        # Create headers for the API request
+        headers = self._create_headers()
+        
+        # Send the request to the API
+        response = requests.post(self.api_url, headers=headers, json=payload, timeout=300)
+        
+        # Handle the response
+        return self._handle_response(response)
+    
+    def __call__(
+        self,
+        messages: List[Dict[str, str]],
+        stop_sequences: Optional[List[str]] = None,
+        grammar: Optional[str] = None,
+        tools_to_call_from: Optional[List[Tool]] = None,
+        custom_role_conversions: Optional[Dict[str, str]] = None,
+        convert_images_to_image_urls: bool = False,
+        **kwargs,
+    ) -> ChatMessage:
+        """
+        Sends a request to the Flow API and returns the response.
+        This method is an alias for generate() for backward compatibility.
+        """
+        return self.generate(
+            messages=messages,
+            stop_sequences=stop_sequences,
+            grammar=grammar,
+            tools_to_call_from=tools_to_call_from,
+            custom_role_conversions=custom_role_conversions,
+            convert_images_to_image_urls=convert_images_to_image_urls,
+            **kwargs,
+        )

cli.py

@@ -0,0 +1,203 @@
+"""Command-line interface for naive_knowledge_base."""
+
+import sys
+import argparse
+from pathlib import Path
+import os
+import dotenv
+
+# Define version locally
+__version__ = "0.1.0"
+
+# Add current directory to path to allow local imports
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# Import dependencies
+try:
+    from smolagents import CodeAgent, ToolCallingAgent
+    from agents.dependency_graph import generate_dependency_graph
+    from api_models import FlowApiModel
+    from tools import (
+        save_content_to_file,
+        read_file_content,
+        delete_folder_or_file,
+        generate_folder_tree,
+    )
+except ImportError as e:
+    print(f"Error importing dependencies: {e}", file=sys.stderr)
+    print("Please install dependencies: pip install -r requirements.txt", file=sys.stderr)
+    sys.exit(1)
+
+
+def run_analysis(source_directory: str, file_extensions: str = "java", 
+                 ignore_dirs: str = "target,.git,test,node_modules,Pods") -> str:
+    """
+    Run dependency graph analysis on a source directory.
+    
+    Args:
+        source_directory: Path to the source code directory to analyze
+        file_extensions: File extensions to analyze (default: "java")
+        ignore_dirs: Comma-separated list of directories to ignore
+        
+    Returns:
+        str: Analysis results
+    """
+    dotenv.load_dotenv()
+    
+    dependency_graph_agent = ToolCallingAgent(
+        tools=[generate_dependency_graph],
+        planning_interval=1,
+        final_answer_checks=[],
+        model=FlowApiModel(
+            model_id="gpt-4.1",
+            temperature=0.4,
+        ),
+        max_steps=7,
+        name="dependency_graph_agent",
+        description="Generates a dependency graph for the given source directory.",
+    )
+    
+    manager_agent = CodeAgent(
+        managed_agents=[dependency_graph_agent],
+        model=FlowApiModel(model_id="gpt-4.1"),
+        tools=[
+            read_file_content,
+            save_content_to_file,
+            delete_folder_or_file,
+            generate_folder_tree
+        ],
+        additional_authorized_imports=[
+            "os",
+            "json",
+            "pandas",
+            "numpy",
+            "pathlib",
+            "ast",
+            "re",
+            "networkx",
+            "collections",
+        ],
+        planning_interval=3,
+        max_steps=30,
+        name="tech_lead_agent",
+        description="This agent is responsible for building and analyzing dependency graphs from source code.",
+    )
+    
+    result = manager_agent.run(
+        f"Generate the tree and dependency graph of the {file_extensions} files in the source: "
+        f"{source_directory}. Make sure to ignore the {ignore_dirs} directories."
+    )
+    
+    return result
+
+
+def create_parser():
+    """Create the argument parser."""
+    parser = argparse.ArgumentParser(
+        prog="naive-kb",
+        description="Analyze code dependencies and generate dependency graphs using AI",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Analyze a Java project
+  naive-kb /path/to/project
+
+  # Analyze with specific file extensions
+  naive-kb /path/to/project --extensions python
+
+  # Ignore specific directories
+  naive-kb /path/to/project --ignore "target,build,.git,node_modules"
+
+  # Full example
+  naive-kb /path/to/project --extensions java --ignore "target,.git,test"
+        """
+    )
+    
+    parser.add_argument(
+        "source_directory",
+        type=str,
+        help="Path to the source code directory to analyze"
+    )
+    
+    parser.add_argument(
+        "-e", "--extensions",
+        type=str,
+        default="java",
+        help="File extensions to analyze (default: java)"
+    )
+    
+    parser.add_argument(
+        "-i", "--ignore",
+        type=str,
+        default="target,.git,test,node_modules,Pods",
+        help="Comma-separated list of directories to ignore (default: target,.git,test,node_modules,Pods)"
+    )
+    
+    parser.add_argument(
+        "-v", "--version",
+        action="version",
+        version=f"%(prog)s {__version__}"
+    )
+    
+    parser.add_argument(
+        "--env-file",
+        type=str,
+        help="Path to .env file (default: .env in current directory)"
+    )
+    
+    return parser
+
+
+def main():
+    """Main CLI entry point."""
+    parser = create_parser()
+    args = parser.parse_args()
+    
+    # Load environment variables
+    if args.env_file:
+        dotenv.load_dotenv(args.env_file)
+    else:
+        dotenv.load_dotenv()
+    
+    # Validate source directory
+    source_path = Path(args.source_directory)
+    if not source_path.exists():
+        print(f"Error: Source directory '{args.source_directory}' does not exist.", file=sys.stderr)
+        sys.exit(1)
+    
+    if not source_path.is_dir():
+        print(f"Error: '{args.source_directory}' is not a directory.", file=sys.stderr)
+        sys.exit(1)
+    
+    print(f"Analyzing source directory: {args.source_directory}")
+    print(f"File extensions: {args.extensions}")
+    print(f"Ignoring directories: {args.ignore}")
+    print("\nRunning analysis...")
+    
+    try:
+        result = run_analysis(
+            source_directory=args.source_directory,
+            file_extensions=args.extensions,
+            ignore_dirs=args.ignore
+        )
+        
+        print("\n" + "=" * 80)
+        print("Analysis Complete!")
+        print("=" * 80)
+        print(result)
+        
+        return 0
+        
+    except KeyboardInterrupt:
+        print("\n\nAnalysis interrupted by user.", file=sys.stderr)
+        return 130
+        
+    except Exception as e:
+        print(f"\nError running analysis: {e}", file=sys.stderr)
+        import traceback
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

main.py

@@ -0,0 +1,68 @@
+from requests import RequestException
+from smolagents import CodeAgent, ToolCallingAgent
+from api_models import FlowApiModel
+from tools import save_content_to_file, read_file_content, delete_folder_or_file, generate_folder_tree
+import dotenv
+import sys
+from agents.dependency_graph import generate_dependency_graph
+dotenv.load_dotenv()
+
+model = FlowApiModel(
+    model_id="gpt-4.1",
+    temperature=0.5,
+)
+
+dependency_graph_agent = ToolCallingAgent(
+    tools=[
+        generate_dependency_graph,
+    ],
+    planning_interval=1,
+    final_answer_checks=[],
+    model=FlowApiModel(
+        model_id="gpt-4.1",
+        temperature=0.4,
+    ),
+    max_steps=7,
+    name="dependency_graph_agent",
+    description="Generates a dependency graph for the given source directory.",
+)
+
+manager_agent = CodeAgent(
+    managed_agents=[dependency_graph_agent],
+    model=FlowApiModel(
+        model_id="gpt-4.1",
+    ),
+    tools=[read_file_content, save_content_to_file, delete_folder_or_file, generate_folder_tree],
+    additional_authorized_imports=[
+        "os",
+        "json",
+        "pandas",
+        "numpy",
+        "pathlib",
+        "ast",
+        "re",
+        "networkx",
+        "collections",
+    ],
+    planning_interval=3,
+    max_steps=30,
+    name="tech_lead_agent",
+    description="This agent is responsible for building and analyzing dependency graphs from source code. It analyzes code files, extracts dependencies, and creates comprehensive dependency graphs with metrics and relationships.",
+)
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python main.py <source_directory>", file=sys.stderr)
+        sys.exit(1)
+    
+    source_directory = sys.argv[1]
+    file_extensions = sys.argv[2] if len(sys.argv) > 2 else "java"
+    ignore_dirs = sys.argv[3] if len(sys.argv) > 3 else "target,.git,test,node_modules,Pods"
+    print("Running main agent...")
+    try:
+        result = manager_agent.run(f"""
+Tasks:
+1. Generate the tree and dependency graph of the {file_extensions} files in the source: {source_directory}. Make sure to ignore the {ignore_dirs} directories.
+        """.strip())
+    except Exception as e:
+        print(f"Error running main agent: {e}", file=sys.stderr)