naas-abi 1.0.0__py3-none-any.whl

naas_abi/workflows/ArtificialAnalysisWorkflow.py

@@ -0,0 +1,337 @@
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from enum import Enum
+from pathlib import Path
+from typing import Annotated, Any, Dict, Set
+
+import requests
+from fastapi import APIRouter
+from langchain_core.tools import BaseTool, StructuredTool
+from naas_abi_core.workflow import Workflow, WorkflowConfiguration
+from naas_abi_core.workflow.workflow import WorkflowParameters
+from pydantic import Field
+
+
+@dataclass
+class ArtificialAnalysisWorkflowConfiguration(WorkflowConfiguration):
+    """Configuration for ArtificialAnalysisWorkflow.
+
+    Attributes:
+        api_key (str): Artificial Analysis API key
+        base_url (str): Base URL for Artificial Analysis API
+    """
+
+    api_key: str
+    base_url: str = "https://artificialanalysis.ai/api/v2"
+
+
+class ArtificialAnalysisWorkflowParameters(WorkflowParameters):
+    """Parameters for ArtificialAnalysisWorkflow execution.
+
+    Attributes:
+        endpoint (str): API endpoint to fetch (llms, text-to-image, etc.)
+        include_categories (bool): Include category breakdowns for media endpoints
+        validate_agents_only (bool): Only generate ontology for modules with active agents
+    """
+
+    endpoint: Annotated[
+        str, Field(description="API endpoint to fetch (llms, text-to-image, etc.)")
+    ] = "llms"
+    include_categories: Annotated[
+        bool, Field(description="Include category breakdowns for media endpoints")
+    ] = False
+    validate_agents_only: Annotated[
+        bool, Field(description="Only generate ontology for modules with active agents")
+    ] = True
+
+
+class ArtificialAnalysisWorkflow(Workflow):
+    """Workflow for fetching and storing Artificial Analysis API data."""
+
+    __configuration: ArtificialAnalysisWorkflowConfiguration
+
+    def __init__(self, configuration: ArtificialAnalysisWorkflowConfiguration):
+        self.__configuration = configuration
+
+    def run_workflow(
+        self, parameters: ArtificialAnalysisWorkflowParameters
+    ) -> Dict[str, Any]:
+        """Execute the workflow to fetch and save Artificial Analysis data.
+
+        Args:
+            parameters: Workflow parameters containing endpoint and options
+
+        Returns:
+            Dict containing operation results and file information
+        """
+        print(f"🚀 Starting Artificial Analysis data fetch for {parameters.endpoint}")
+
+        # Step 1: Get valid modules with agents
+        valid_modules = (
+            self._get_modules_with_agents()
+            if parameters.validate_agents_only
+            else set()
+        )
+
+        if parameters.validate_agents_only:
+            print(
+                f"🤖 Found {len(valid_modules)} modules with active agents: {', '.join(sorted(valid_modules))}"
+            )
+
+        # Step 2: Fetch data from API
+        api_data = self._fetch_models_data(
+            parameters.endpoint, parameters.include_categories
+        )
+
+        if api_data.get("status") == "error":
+            return {"status": "error", "message": "Failed to fetch data from API"}
+
+        # Step 3: Filter data for valid modules only
+        filtered_data = api_data
+        if parameters.validate_agents_only and valid_modules:
+            filtered_data = self._filter_data_for_valid_modules(api_data, valid_modules)
+
+        # Step 4: Save filtered data
+        timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
+        filename = f"{timestamp}_{parameters.endpoint}_data.json"
+
+        # Create storage directory
+        storage_dir = Path(
+            "storage/datastore/core/modules/abi/ArtificialAnalysisWorkflow"
+        )
+        storage_dir.mkdir(parents=True, exist_ok=True)
+
+        # Save filtered API response
+        output_file = storage_dir / filename
+        with open(output_file, "w", encoding="utf-8") as f:
+            json.dump(filtered_data, f, indent=2, ensure_ascii=False)
+
+        models_count = len(filtered_data.get("data", []))
+        original_count = len(api_data.get("data", []))
+
+        print(f"💾 Saved {models_count}/{original_count} models to: {output_file}")
+        print(f"📊 Endpoint: {parameters.endpoint}")
+        print(f"🕐 Timestamp: {timestamp}")
+
+        if parameters.validate_agents_only:
+            print("✅ Filtered to models with corresponding agents only")
+
+        return {
+            "status": "success",
+            "endpoint": parameters.endpoint,
+            "models_count": models_count,
+            "original_count": original_count,
+            "valid_modules": list(valid_modules) if valid_modules else [],
+            "output_file": str(output_file),
+            "timestamp": timestamp,
+        }
+
+    def _fetch_models_data(
+        self, endpoint: str, include_categories: bool = False
+    ) -> Dict[str, Any]:
+        """Fetch model data from Artificial Analysis API."""
+        # Different URL structure for LLMs vs media endpoints
+        if endpoint == "llms":
+            url = f"{self.__configuration.base_url}/data/llms/models"
+        else:
+            # Media endpoints like text-to-image, text-to-speech, etc.
+            url = f"{self.__configuration.base_url}/data/media/{endpoint}"
+
+        headers = {
+            "x-api-key": self.__configuration.api_key,
+            "Content-Type": "application/json",
+        }
+
+        params = {}
+        if include_categories:
+            params["include_categories"] = "true"
+
+        try:
+            print(f"🔍 Fetching data from: {url}")
+            response = requests.get(url, headers=headers, params=params)
+            response.raise_for_status()
+
+            data = response.json()
+            models_count = len(data.get("data", []))
+            print(f"✅ Successfully fetched {models_count} models from {endpoint}")
+
+            return {
+                "status": "success",
+                "data": data.get("data", []),
+                "metadata": {
+                    "endpoint": endpoint,
+                    "total_models": models_count,
+                    "fetched_at": datetime.now(timezone.utc).isoformat(),
+                    "api_url": url,
+                },
+            }
+
+        except requests.exceptions.RequestException as e:
+            print(f"❌ Error fetching data: {e}")
+            return {"status": "error", "error": str(e), "endpoint": endpoint}
+
+    def _get_modules_with_agents(self) -> Set[str]:
+        """Get set of module names that have agents folders with *Agent.py files.
+
+        Returns:
+            Set[str]: Module names that have active agents
+        """
+        valid_modules: Set[str] = set()
+        core_modules_path = Path("src/core/modules")
+
+        if not core_modules_path.exists():
+            print(f"⚠️  Core modules path not found: {core_modules_path}")
+            return valid_modules
+
+        print(f"🔍 Scanning for modules with agents in: {core_modules_path}")
+
+        # Scan each module directory
+        for module_dir in core_modules_path.iterdir():
+            if not module_dir.is_dir() or module_dir.name.startswith("."):
+                continue
+
+            agents_dir = module_dir / "agents"
+            if not agents_dir.exists():
+                continue
+
+            # Look for *Agent.py files
+            agent_files = list(agents_dir.glob("*Agent.py"))
+            if agent_files:
+                valid_modules.add(module_dir.name)
+                print(f"  ✅ {module_dir.name}: Found {len(agent_files)} agent files")
+            else:
+                print(
+                    f"  ❌ {module_dir.name}: Has agents/ folder but no *Agent.py files"
+                )
+
+        return valid_modules
+
+    def _filter_data_for_valid_modules(
+        self, api_data: Dict[str, Any], valid_modules: Set[str]
+    ) -> Dict[str, Any]:
+        """Filter API data to only include models that have corresponding agent modules.
+
+        Args:
+            api_data: Raw API data from Artificial Analysis
+            valid_modules: Set of module names with agents
+
+        Returns:
+            Dict[str, Any]: Filtered API data
+        """
+        if not api_data.get("data"):
+            return api_data
+
+        filtered_models = []
+        skipped_models = []
+
+        for model in api_data["data"]:
+            # Extract provider/module name from model data
+            module_name = self._extract_module_name_from_model(model)
+
+            if module_name in valid_modules:
+                filtered_models.append(model)
+            else:
+                skipped_models.append(
+                    {
+                        "name": model.get("name", "Unknown"),
+                        "provider": model.get("provider", "Unknown"),
+                        "reason": f"No agent found for module: {module_name}",
+                    }
+                )
+
+        print("📋 Filtered results:")
+        print(f"  ✅ Keeping {len(filtered_models)} models with agents")
+        print(f"  ❌ Skipping {len(skipped_models)} models without agents")
+
+        if skipped_models:
+            print("  📝 Skipped models:")
+            for skipped in skipped_models[:5]:  # Show first 5
+                print(
+                    f"    • {skipped['name']} ({skipped['provider']}) - {skipped['reason']}"
+                )
+            if len(skipped_models) > 5:
+                print(f"    • ... and {len(skipped_models) - 5} more")
+
+        # Return filtered data with same structure
+        filtered_data = api_data.copy()
+        filtered_data["data"] = filtered_models
+
+        return filtered_data
+
+    def _extract_module_name_from_model(self, model: Dict[str, Any]) -> str:
+        """Extract the module name from a model's provider information.
+
+        Args:
+            model: Model data from API
+
+        Returns:
+            str: Module name (e.g., 'openai', 'anthropic', 'mistral')
+        """
+        provider = model.get("provider", "").lower()
+        model_name = model.get("name", "").lower()
+
+        # Map providers to module names
+        provider_mapping = {
+            "openai": "chatgpt",
+            "anthropic": "claude",
+            "mistral": "mistral",
+            "google": "gemini",
+            "meta": "llama",
+            "perplexity": "perplexity",
+            "xai": "grok",
+            "together": "together",
+            "groq": "groq",
+        }
+
+        # Try to match by provider first
+        for provider_key, module_name in provider_mapping.items():
+            if provider_key in provider:
+                return module_name
+
+        # Try to match by model name patterns
+        if "gpt" in model_name or "openai" in model_name:
+            return "chatgpt"
+        elif "claude" in model_name or "anthropic" in model_name:
+            return "claude"
+        elif "mistral" in model_name:
+            return "mistral"
+        elif "gemini" in model_name or "palm" in model_name:
+            return "gemini"
+        elif "llama" in model_name or "meta" in model_name:
+            return "llama"
+        elif "grok" in model_name or "xai" in model_name:
+            return "grok"
+
+        # Default to provider name if no mapping found
+        return provider.replace(" ", "").replace("-", "").lower()
+
+    def as_tools(self) -> list[BaseTool]:
+        """Returns a list of LangChain tools for this workflow.
+
+        Returns:
+            list[BaseTool]: List containing the workflow tool
+        """
+        return [
+            StructuredTool(
+                name="artificial_analysis_data_fetch",
+                description="Fetch AI model data from Artificial Analysis API, filter for modules with active agents, and save as timestamped JSON files",
+                func=lambda **kwargs: self.run(
+                    ArtificialAnalysisWorkflowParameters(**kwargs)
+                ),
+                args_schema=ArtificialAnalysisWorkflowParameters,
+            )
+        ]
+
+    def as_api(
+        self,
+        router: APIRouter,
+        route_name: str = "",
+        name: str = "",
+        description: str = "",
+        description_stream: str = "",
+        tags: list[str | Enum] | None = None,
+    ) -> None:
+        if tags is None:
+            tags = []
+        return None

naas_abi/workflows/ArtificialAnalysisWorkflow_test.py

@@ -0,0 +1,57 @@
+import os
+from pathlib import Path
+
+import pytest
+from naas_abi.workflows.ArtificialAnalysisWorkflow import (
+    ArtificialAnalysisWorkflow,
+    ArtificialAnalysisWorkflowConfiguration,
+    ArtificialAnalysisWorkflowParameters,
+)
+
+
+@pytest.fixture
+def workflow() -> ArtificialAnalysisWorkflow:
+    api_key = os.getenv("AA_AI_API_KEY", "test_key")
+
+    workflow_configuration = ArtificialAnalysisWorkflowConfiguration(
+        api_key=api_key, base_url="https://artificialanalysis.ai/api/v2"
+    )
+
+    return ArtificialAnalysisWorkflow(workflow_configuration)
+
+
+@pytest.mark.skipif(
+    not os.getenv("AA_AI_API_KEY"), reason="AA_AI_API_KEY environment variable not set"
+)
+def test_workflow_fetch_llms_data(workflow: ArtificialAnalysisWorkflow):
+    """Test fetching LLMs data from Artificial Analysis API."""
+    result = workflow.run_workflow(
+        ArtificialAnalysisWorkflowParameters(endpoint="llms", include_categories=False)
+    )
+
+    assert result is not None, result
+    assert result["status"] == "success", result
+    assert result["models_count"] > 0, result
+    assert "output_file" in result, result
+    assert "timestamp" in result, result
+
+    # Verify file was created
+    output_file = Path(result["output_file"])
+    assert output_file.exists(), f"Output file {output_file} was not created"
+
+    # Verify file contains valid JSON
+    import json
+
+    with open(output_file, "r") as f:
+        data = json.load(f)
+    assert "data" in data, "JSON file should contain 'data' field"
+    assert "metadata" in data, "JSON file should contain 'metadata' field"
+
+
+def test_workflow_as_tools(workflow: ArtificialAnalysisWorkflow):
+    """Test that workflow can be converted to LangChain tools."""
+    tools = workflow.as_tools()
+
+    assert len(tools) == 1, f"Expected 1 tool, got {len(tools)}"
+    assert tools[0].name == "artificial_analysis_data_fetch"
+    assert "Artificial Analysis" in tools[0].description

naas_abi/workflows/ConvertOntologyGraphToYamlWorkflow.py

@@ -0,0 +1,210 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Annotated, Dict, Optional
+
+import pydash as _
+import yaml
+from fastapi import APIRouter
+from langchain_core.tools import BaseTool, StructuredTool
+from naas_abi import config, logger
+from naas_abi.mappings import COLORS_NODES
+from naas_abi_core.utils.OntologyYaml import OntologyYaml
+from naas_abi_core.workflow import Workflow, WorkflowConfiguration
+from naas_abi_core.workflow.workflow import WorkflowParameters
+from naas_abi_marketplace.applications.naas.integrations.NaasIntegration import (
+    NaasIntegration,
+    NaasIntegrationConfiguration,
+)
+from pydantic import Field
+from rdflib import Graph
+from yaml import Dumper
+
+
+@dataclass
+class ConvertOntologyGraphToYamlWorkflowConfiguration(WorkflowConfiguration):
+    """Configuration for ConvertOntologyGraphToYaml workflow.
+
+    Attributes:
+        naas_integration_config (NaasIntegrationConfiguration): Configuration for the Naas integration
+    """
+
+    naas_integration_config: NaasIntegrationConfiguration
+
+
+class ConvertOntologyGraphToYamlWorkflowParameters(WorkflowParameters):
+    """Parameters for ConvertOntologyGraphToYaml workflow execution.
+
+    Attributes:
+        ontology_name (str): The name of the ontology store to use
+        label (str): The label of the ontology
+        description (str): The description of the ontology
+        logo_url (str): The URL of the ontology logo
+        level (str): The level of the ontology (e.g., 'TOP_LEVEL', 'MID_LEVEL', 'DOMAIN', 'USE_CASE')
+        display_relations_names (bool): Whether to display relation names in the visualization
+    """
+
+    graph: Annotated[
+        str, Field(..., description="The graph serialized as turtle format")
+    ]
+    ontology_id: Annotated[
+        Optional[str], Field(..., description="The ID of the ontology")
+    ] = None
+    label: Annotated[str, Field(..., description="The label of the ontology")] = (
+        "New Ontology"
+    )
+    description: Annotated[
+        str,
+        Field(
+            ...,
+            description="The description of the ontology. Example: 'Represents ABI Ontology with agents, workflows, ontologies, pipelines and integrations.'",
+        ),
+    ] = "New Ontology Description"
+    logo_url: Annotated[
+        Optional[str], Field(..., description="The logo URL of the ontology")
+    ] = "https://naasai-public.s3.eu-west-3.amazonaws.com/abi-demo/ontology_ULO.png"
+    level: Annotated[str, Field(..., description="The level of the ontology")] = (
+        "USE_CASE"
+    )
+    display_relations_names: Annotated[
+        bool,
+        Field(
+            ..., description="Whether to display relation names in the visualization"
+        ),
+    ] = True
+    class_colors_mapping: Annotated[
+        Dict, Field(..., description="The mapping of class colors")
+    ] = COLORS_NODES
+
+
+class ConvertOntologyGraphToYamlWorkflow(Workflow):
+    """Workflow for converting ontology graph to YAML."""
+
+    __configuration: ConvertOntologyGraphToYamlWorkflowConfiguration
+
+    def __init__(self, configuration: ConvertOntologyGraphToYamlWorkflowConfiguration):
+        self.__configuration = configuration
+        self.__naas_integration = NaasIntegration(
+            self.__configuration.naas_integration_config
+        )
+
+    def graph_to_yaml(
+        self, parameters: ConvertOntologyGraphToYamlWorkflowParameters
+    ) -> str:
+        # Initialize parameters
+        logger.debug(f"==> Converting ontology graph to YAML: {parameters.label}")
+        yaml_data = None
+        ontology_id = parameters.ontology_id
+
+        # Create Graph from turtle string
+        g = Graph()
+        g.parse(data=parameters.graph, format="turtle")
+
+        # Upload asset to Naas
+        asset = self.__naas_integration.upload_asset(
+            data=parameters.graph.encode("utf-8"),  # Use the original turtle string
+            workspace_id=config.workspace_id,
+            storage_name=config.storage_name,
+            prefix="assets",
+            object_name=str(parameters.label + ".ttl"),
+            visibility="public",
+        )
+        # Save asset URL to JSON
+        if asset is None:
+            raise ValueError("Failed to upload asset to Naas")
+
+        asset_url = asset.get("asset", {}).get("url")
+        if not asset_url:
+            raise ValueError("Asset URL not found in response")
+
+        if asset_url.endswith("/"):
+            asset_url = asset_url[:-1]
+
+        # Convert to YAML
+        try:
+            yaml_data = OntologyYaml.rdf_to_yaml(
+                g,
+                display_relations_names=parameters.display_relations_names,
+                class_colors_mapping=parameters.class_colors_mapping,
+            )
+        except Exception as e:
+            message = f"Error converting ontology to YAML: {e}"
+            raise e
+
+        # Initialize parameters
+        if yaml_data is not None:
+            workspace_id = config.workspace_id
+            onto_label = parameters.label
+            onto_description = parameters.description
+            onto_logo_url = parameters.logo_url
+            onto_level = parameters.level
+
+            # Get ontology ID if it exists
+            ontologies = self.__naas_integration.get_ontologies(workspace_id).get(
+                "ontologies", []
+            )
+            for ontology in ontologies:
+                if (
+                    ontology
+                    and isinstance(ontology, dict)
+                    and ontology.get("label") == onto_label
+                ):
+                    ontology_id = ontology.get("id")
+                    break
+
+            if ontology_id is None:
+                # Create new ontology
+                res = self.__naas_integration.create_ontology(
+                    workspace_id=workspace_id,
+                    label=onto_label,
+                    source=yaml.dump(yaml_data, Dumper=Dumper),
+                    level=onto_level,
+                    description=onto_description,
+                    download_url=asset_url,
+                    logo_url=onto_logo_url,
+                )
+                ontology_id = str(_.get(res, "ontology.id", ""))
+                message = (
+                    f"✅ Ontology '{onto_label}' ({ontology_id}) successfully created."
+                )
+            else:
+                # Update existing ontology
+                res = self.__naas_integration.update_ontology(
+                    workspace_id=workspace_id,
+                    ontology_id=ontology_id,
+                    source=yaml.dump(yaml_data, Dumper=Dumper),
+                    level=onto_level,
+                    description=onto_description,
+                    download_url=asset_url,
+                    logo_url=onto_logo_url,
+                )
+                message = (
+                    f"✅ Ontology '{onto_label}' ({ontology_id}) successfully updated."
+                )
+        logger.info(message)
+        if ontology_id is None:
+            raise ValueError("Failed to create or update ontology")
+        return ontology_id
+
+    def as_tools(self) -> list[BaseTool]:
+        """Returns a list of LangChain tools for this workflow."""
+        return [
+            StructuredTool(
+                name="convert_graph_to_yaml",
+                description="Convert an ontology graph to YAML.",
+                func=lambda **kwargs: self.graph_to_yaml(
+                    ConvertOntologyGraphToYamlWorkflowParameters(**kwargs)
+                ),
+                args_schema=ConvertOntologyGraphToYamlWorkflowParameters,
+            )
+        ]
+
+    def as_api(
+        self,
+        router: APIRouter,
+        route_name: str = "",
+        name: str = "",
+        description: str = "",
+        description_stream: str = "",
+        tags: list[str | Enum] | None = None,
+    ) -> None:
+        pass

naas_abi/workflows/ConvertOntologyGraphToYamlWorkflow_test.py

@@ -0,0 +1,78 @@
+import pytest
+from naas_abi import secret
+from naas_abi.workflows.ConvertOntologyGraphToYamlWorkflow import (
+    ConvertOntologyGraphToYamlWorkflow,
+    ConvertOntologyGraphToYamlWorkflowConfiguration,
+    ConvertOntologyGraphToYamlWorkflowParameters,
+)
+from naas_abi_marketplace.applications.naas.integrations.NaasIntegration import (
+    NaasIntegration,
+    NaasIntegrationConfiguration,
+)
+
+
+@pytest.fixture
+def workflow() -> ConvertOntologyGraphToYamlWorkflow:
+    return ConvertOntologyGraphToYamlWorkflow(
+        ConvertOntologyGraphToYamlWorkflowConfiguration(
+            NaasIntegrationConfiguration(api_key=secret.get("NAAS_API_KEY"))
+        )
+    )
+
+
+def test_convert_ontology_graph_to_yaml_workflow(
+    workflow: ConvertOntologyGraphToYamlWorkflow,
+):
+    import time
+    from uuid import uuid4
+
+    from naas_abi import config, logger, services
+    from rdflib import OWL, RDF, RDFS, Graph, Literal, Namespace, URIRef
+
+    ABI = Namespace("http://ontology.naas.ai/abi/")
+
+    graph = Graph()
+    uri = ABI[str(uuid4())]
+    graph.add(
+        (uri, RDF.type, URIRef("https://www.commoncoreontologies.org/ont00000443"))
+    )
+    graph.add((uri, RDF.type, OWL.NamedIndividual))
+    graph.add((uri, RDFS.label, Literal("Naas.ai")))
+    graph.add(
+        (
+            uri,
+            ABI.logo,
+            Literal(
+                "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcQ9gXMaBLQZ39W6Pk53PRuzFjUvv_6lLRWPoQ&s"
+            ),
+        )
+    )
+    services.triple_store_service.insert(graph)
+    time.sleep(3)
+
+    # Run workflow
+    ontology_id = workflow.graph_to_yaml(
+        ConvertOntologyGraphToYamlWorkflowParameters(
+            graph=graph.serialize(format="turtle"),
+            ontology_id=None,
+            label="Naas.ai",
+            description="Naas.ai Ontology",
+            logo_url="https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcQ9gXMaBLQZ39W6Pk53PRuzFjUvv_6lLRWPoQ&s",
+            level="USE_CASE",
+            display_relations_names=True,
+            class_colors_mapping={},
+        )
+    )
+
+    # Remove graph
+    services.triple_store_service.remove(graph)
+
+    # Remove ontology
+    naas_integration = NaasIntegration(
+        NaasIntegrationConfiguration(api_key=secret.get("NAAS_API_KEY"))
+    )
+    result = naas_integration.delete_ontology(
+        workspace_id=config.workspace_id, ontology_id=ontology_id
+    )
+    logger.info(f"Removed ontology: {result}")
+    assert result is not None, result