ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/tools/pricestool.py

@@ -0,0 +1,146 @@
+from typing import List
+from pydantic import BaseModel, Field
+from datetime import date
+from parrot.tools.querytoolkit import QueryToolkit
+from parrot.tools.decorators import tool_schema
+
+
+class PriceInput(BaseModel):
+    tenant: str = Field(..., description="Program or tenant identifier (e.g. hisense, epson, etc).")
+    model: str = Field(..., description="Model number of the product to query the price for.")
+    week: int = Field(..., description="Week number for which to retrieve the price.")
+    output_format: str = Field("structured", description="Output format: 'string' or 'structured'.")
+
+
+class ModelPriceInput(BaseModel):
+    tenant: str = Field(..., description="Program or tenant identifier (e.g. hisense, epson, etc).")
+    model: str = Field(..., description="Model number of the product to query the price for.")
+    output_format: str = Field(default="pandas", description="Output format: 'string' or 'structured'.")
+    limit: int = Field(10, description="Number of records to retrieve.")
+
+
+class WeeklyPriceInput(BaseModel):
+    tenant: str = Field(..., description="Program or tenant identifier (e.g. hisense, epson, etc).")
+    week: int = Field(..., description="Week number for which to retrieve the price.")
+    output_format: str = Field(default="pandas", description="Output format: 'string' or 'structured'.")
+    limit: int = Field(10, description="Number of records to retrieve.")
+
+class TotalPriceInput(BaseModel):
+    tenant: str = Field(..., description="Program or tenant identifier (e.g. hisense, epson, etc).")
+    start_date: str = Field(..., description="Start date for the price query (YYYY-MM-DD).")
+    end_date: str = Field(..., description="End date for the price query (YYYY-MM-DD).")
+    output_format: str = Field(default="pandas", description="Output format: 'string' or 'structured'.")
+    limit: int = Field(50, description="Number of records to retrieve.")
+
+
+class PriceOutput(BaseModel):
+    product_id: str = Field(..., description="Unique identifier for the product.")
+    price: float = Field(..., description="Price of the product.")
+    week: int = Field(..., description="Week number for which the price is applicable.")
+    start_date: date = Field(..., description="Start date of the pricing period.")
+    end_date: date = Field(..., description="End date of the pricing period.")
+
+
+
+class PricesTool(QueryToolkit):
+    """Tool for querying product prices from a database or API."""
+
+    name = "prices_tool"
+    description = "A tool to query product prices."
+
+
+    @tool_schema(PriceInput, description="The price information as a string.")
+    async def get_model_price(
+        self,
+        tenant: str,
+        model: str,
+        week: int,
+        output_format: str = "structured",
+    ) -> PriceOutput:
+        """Fetches the price of a product for a given tenant, model, and week.
+
+        Args:
+            tenant (str): The program or tenant identifier.
+            model (str): The model number of the product.
+            week (int): The week number for     which to retrieve the price.
+        """
+        sql = await self._get_query("get_pricing")
+        sql = sql.format(
+            tenant=tenant,
+            model=model,
+            week=week
+        )
+        try:
+            return await self._get_dataset(
+                sql,
+                output_format=output_format,
+                structured_obj=PriceOutput if output_format == "structured" else None
+            )
+        except ValueError as ve:
+            return f"No Pricing data found for the specified tenant and product, error: {ve}"
+        except Exception as e:
+            return f"Error fetching pricing data: {e}"
+
+    @tool_schema(WeeklyPriceInput, description="The price information as a string.")
+    async def get_weekly_price(
+        self,
+        tenant: str,
+        week: int,
+        limit: int = 10,
+        output_format: str = "structured",
+    ) -> List[PriceOutput]:
+        """Fetches all product prices for a given tenant and week.
+
+        Args:
+            tenant (str): The program or tenant identifier.
+            week (int): The week number for which to retrieve the price.
+        """
+        sql = await self._get_query("get_weekly_pricing")
+        sql = sql.format(
+            tenant=tenant,
+            week=week,
+            limit=limit
+        )
+        try:
+            return await self._get_dataset(
+                sql,
+                output_format=output_format,
+                structured_obj=PriceOutput if output_format == "structured" else None
+            )
+        except ValueError as ve:
+            return f"No Pricing data found for the specified tenant and product, error: {ve}"
+        except Exception as e:
+            return f"Error fetching pricing data: {e}"
+        
+    @tool_schema(TotalPriceInput, description="The price information as a string.")
+    async def get_price(
+        self,
+        tenant: str,
+        start_date: str,
+        end_date: str,
+        limit: int = 10,
+        output_format: str = "structured",
+    ) -> List[PriceOutput]:
+        """Fetches all product prices for a given tenant and date range.
+
+        Args:
+            tenant (str): The program or tenant identifier.
+            week (int): The week number for which to retrieve the price.
+        """
+        sql = await self._get_query("get_total_pricing")
+        sql = sql.format(
+            tenant=tenant,
+            start_date=start_date,
+            end_date=end_date,
+            limit=limit
+        )
+        try:
+            return await self._get_dataset(
+                sql,
+                output_format=output_format,
+                structured_obj=PriceOutput if output_format == "structured" else None
+            )
+        except ValueError as ve:
+            return f"No Pricing data found for the specified tenant and product, error: {ve}"
+        except Exception as e:
+            return f"Error fetching pricing data: {e}"

parrot/tools/products/__init__.py

@@ -0,0 +1,246 @@
+from typing import Optional, Dict, Union, List, Any
+from datetime import datetime
+import json
+from pydantic import BaseModel, Field, field_validator, ConfigDict
+from navconfig import BASE_DIR
+from asyncdb import AsyncDB
+from asyncdb.models import Model, Field as ModelField
+from querysource.conf import default_dsn
+from ..abstract import AbstractTool
+
+
+
+class ProductInput(BaseModel):
+    """Input schema for product information requests."""
+    model: str = Field(
+        ..., description="The product model identifier (e.g., 'X1234', 'Y5678')."
+    )
+    program_slug: str = Field(
+        ..., description="The program slug associated with the product (e.g., 'alpha', 'beta')."
+    )
+
+
+class ProductInfo(BaseModel):
+    """Schema for the product information returned by the query."""
+    name: str
+    model: str
+    description: str
+    picture_url: str
+    brand: str
+    # pricing: Decimal
+    pricing: str
+    customer_satisfaction: Optional[str] = None
+    product_evaluation: Optional[str] = None
+    product_compliant: Optional[str] = None
+    # specifications: Dict[str, Union[dict, list]] = Field(
+    #     default_factory=dict,
+    #     description="Specifications of the product, can be a dict or list."
+    # )
+    specifications: Dict[str, Union[str, int, float, bool, list, dict]] = Field(
+        default_factory=dict,
+        description="Specifications of the product as a dictionary."
+    )
+    review_average: float
+    reviews: int
+
+    @field_validator('specifications', mode='before')
+    @classmethod
+    def parse_specifications(cls, v):
+        if v is None or v == '':
+            return {}
+        if isinstance(v, dict):
+            return v
+        if isinstance(v, (bytes, bytearray)):
+            v = v.decode('utf-8', errors='ignore')
+        if isinstance(v, str):
+            try:
+                parsed = json.loads(v)
+            except json.JSONDecodeError as e:
+                raise ValueError("Specifications field is not a valid JSON string.") from e
+            if not isinstance(parsed, dict):
+                raise TypeError("Specifications JSON must decode to a dictionary.")
+            return parsed
+        raise TypeError("specifications must be a dict or a JSON string.")
+
+    # Add a model_config to prevent additional properties
+    model_config = ConfigDict(
+        arbitrary_types_allowed=False,
+        extra="forbid",
+    )
+
+class ProductInfoTool(AbstractTool):
+    """Tool to get detailed information about a specific product model."""
+    name = "get_product_information"
+    description = (
+        "Use this tool to get detailed information about a specific product model. "
+        "Provide the exact model identifier as input."
+    )
+    args_schema = ProductInput
+
+    async def _execute(self, model: str, program_slug: str) -> ProductInfo:
+        db = AsyncDB('pg', dsn=default_dsn)
+
+        # Use static_dir if configured, otherwise fall back to BASE_DIR
+        base_path = self.static_dir if hasattr(self, 'static_dir') and self.static_dir else BASE_DIR
+
+        # Try multiple paths for backward compatibility
+        # 1. Direct path (for when static_dir points to programs/ or taskstore/programs/)
+        query_file = base_path / program_slug / 'sql' / 'products.sql'
+
+        # 2. Try with 'programs/' prefix (for when static_dir points to base directory)
+        if not query_file.exists():
+            query_file = base_path / 'programs' / program_slug / 'sql' / 'products.sql'
+
+        # 3. Fallback to old structure: agents/product_report/{program_slug}/products.sql
+        if not query_file.exists():
+            query_file = base_path / 'agents' / 'product_report' / program_slug / 'products.sql'
+
+        if not query_file.exists():
+            raise FileNotFoundError(
+                f"Query file not found for program_slug '{program_slug}'. Tried:\n"
+                f"  - {base_path / program_slug / 'sql' / 'products.sql'}\n"
+                f"  - {base_path / 'programs' / program_slug / 'sql' / 'products.sql'}\n"
+                f"  - {base_path / 'agents' / 'product_report' / program_slug / 'products.sql'}"
+            )
+
+        query = query_file.read_text()
+        async with await db.connection() as conn:  # noqa
+            product_data, error = await conn.query(query, model)
+            if error:
+                raise RuntimeError(f"Database query failed: {error}")
+            if not product_data:
+                raise ValueError(f"No product found with model '{model}' in program '{program_slug}'.")
+
+            return ProductInfo(**product_data[0])
+
+
+class ProductListInput(BaseModel):
+    """Input schema for product list requests."""
+    program_slug: str = Field(
+        ..., description="The program slug to get products from (e.g., 'google', 'hisense')."
+    )
+    models: Optional[List[str]] = Field(
+        default=None, description="Optional list of specific models to get. If None, gets all models."
+    )
+
+
+class ProductListTool(AbstractTool):
+    """Tool to get list of products for a given program/tenant."""
+    name = "get_products_list"
+    description = (
+        "Use this tool to get a list of products for a given program/tenant. "
+        "Provide the program slug as input. Optionally provide a list of specific models."
+    )
+    args_schema = ProductListInput
+
+    async def _execute(self, program_slug: str, models: Optional[List[str]] = None) -> List[Dict[str, str]]:
+        """Get list of products for a program."""
+        db = AsyncDB('pg', dsn=default_dsn)
+
+        # Use static_dir if configured, otherwise fall back to BASE_DIR
+        base_path = self.static_dir if hasattr(self, 'static_dir') and self.static_dir else BASE_DIR
+
+        # Determine which SQL file to use
+        file_type = 'product_single.sql' if models else 'products_list.sql'
+
+        # Try multiple paths for backward compatibility
+        # 1. Direct path (for when static_dir points to programs/ or taskstore/programs/)
+        query_file = base_path / program_slug / 'sql' / file_type
+
+        # 2. Try with 'programs/' prefix (for when static_dir points to base directory)
+        if not query_file.exists():
+            query_file = base_path / 'programs' / program_slug / 'sql' / file_type
+
+        # 3. Fallback to old structure: agents/product_report/{program_slug}/<file>
+        if not query_file.exists():
+            query_file = base_path / 'agents' / 'product_report' / program_slug / file_type
+
+        if not query_file.exists():
+            raise FileNotFoundError(
+                f"Products query file not found for program_slug '{program_slug}'. Tried:\n"
+                f"  - {base_path / program_slug / 'sql' / file_type}\n"
+                f"  - {base_path / 'programs' / program_slug / 'sql' / file_type}\n"
+                f"  - {base_path / 'agents' / 'product_report' / program_slug / file_type}"
+            )
+
+        query = query_file.read_text()
+        async with await db.connection() as conn:  # noqa
+            if models:
+                # Execute with models parameter
+                products, error = await conn.query(query, models)
+            else:
+                # Execute without parameters
+                products, error = await conn.query(query)
+
+            if error:
+                raise RuntimeError(f"Database query failed: {error}")
+            if not products:
+                return []
+
+            return products
+
+
+class ProductResponse(Model):
+    """
+    ProductResponse is a model that defines the structure of the response for Product agents.
+    """
+    model: Optional[str] = ModelField(
+        default=None,
+        description="Model of the product"
+    )
+    program_slug: Optional[str] = ModelField(
+        default=None,
+        description="Program/tenant identifier"
+    )
+    agent_id: Optional[str] = ModelField(
+        default=None,
+        description="Unique identifier for the agent that processed the request"
+    )
+    agent_name: Optional[str] = ModelField(
+        default="ProductReport",
+        description="Name of the agent that processed the request"
+    )
+    status: str = ModelField(default="success", description="Status of the response")
+    data: Optional[str] = ModelField(
+        default=None,
+        description="Data returned by the agent, can be text, JSON, etc."
+    )
+    # Optional output field for structured data
+    output: Optional[Any] = ModelField(
+        default=None,
+        description="Output of the agent's processing"
+    )
+    attributes: Dict[str, str] = ModelField(
+        default_factory=dict,
+        description="Attributes associated with the response"
+    )
+    # Timestamp
+    created_at: datetime = ModelField(
+        default_factory=datetime.now, description="Timestamp when response was created"
+    )
+    # Optional file paths
+    transcript: Optional[str] = ModelField(
+        default=None, description="Transcript of the conversation with the agent"
+    )
+    script_path: Optional[str] = ModelField(
+        default=None, description="Path to the conversational script associated with the session"
+    )
+    podcast_path: Optional[str] = ModelField(
+        default=None, description="Path to the podcast associated with the session"
+    )
+    pdf_path: Optional[str] = ModelField(
+        default=None, description="Path to the PDF associated with the session"
+    )
+    document_path: Optional[str] = ModelField(
+        default=None, description="Path to any document generated during session"
+    )
+    # complete list of generated files:
+    files: List[str] = ModelField(
+        default_factory=list, description="List of documents generated during the session")
+
+    class Meta:
+        """Meta class for ProductResponse."""
+        name = "products_informations"
+        schema = "product_report"
+        strict = True
+        frozen = False

parrot/tools/prophet_tool.py

@@ -0,0 +1,171 @@
+"""
+ProphetForecastTool for time series forecasting using Facebook Prophet.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import pandas as pd
+from pydantic import Field, field_validator
+from prophet import Prophet
+
+from .abstract import AbstractTool, AbstractToolArgsSchema, ToolResult
+
+
+class ProphetForecastArgs(AbstractToolArgsSchema):
+    """Arguments for :class:`ProphetForecastTool`."""
+
+    dataframe: str = Field(
+        ..., description="Name or alias of the DataFrame containing the time series"
+    )
+    ds_column: str = Field(
+        ..., description="Column with datestamp values (will be converted to datetime)"
+    )
+    y_column: str = Field(
+        ..., description="Numeric column to forecast"
+    )
+    periods: int = Field(
+        365,
+        description="Number of future periods to forecast",
+        ge=1,
+    )
+    freq: str = Field(
+        "D",
+        description="Pandas frequency string for future periods (e.g., 'D', 'W', 'M')",
+        min_length=1,
+    )
+    include_history: bool = Field(
+        default=True,
+        description="Whether to include the historical dates in the forecast output",
+    )
+
+    @field_validator("periods")
+    @classmethod
+    def validate_periods(cls, value: int) -> int:
+        if value < 1:
+            raise ValueError("periods must be a positive integer")
+        return value
+
+
+class ProphetForecastTool(AbstractTool):
+    """Generate time series forecasts with Facebook Prophet and return plots."""
+
+    name = "prophet_forecast"
+    description = (
+        "Fit a Facebook Prophet model on a DataFrame and generate future forecasts "
+        "with corresponding forecast plot image."
+    )
+    args_schema = ProphetForecastArgs
+
+    def __init__(
+        self,
+        dataframes: Optional[Dict[str, pd.DataFrame]] = None,
+        alias_map: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.dataframes: Dict[str, pd.DataFrame] = dataframes or {}
+        self.alias_map: Dict[str, str] = alias_map or {}
+
+    def _default_output_dir(self) -> Path:
+        return self.static_dir / "reports" / "prophet_forecast"
+
+    def update_context(
+        self, dataframes: Dict[str, pd.DataFrame], alias_map: Optional[Dict[str, str]] = None
+    ) -> None:
+        """Update internal references to available DataFrames and aliases."""
+
+        self.dataframes = dataframes
+        if alias_map is not None:
+            self.alias_map = alias_map
+
+    async def _execute(self, **kwargs: Any) -> ToolResult:
+        args = self.args_schema(**kwargs)
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self._run_forecast, args)
+
+    def _run_forecast(self, args: ProphetForecastArgs) -> ToolResult:
+        dataframe = self._resolve_dataframe(args.dataframe)
+        cleaned_df = self._prepare_dataframe(dataframe, args)
+
+        model = Prophet()
+        model.fit(cleaned_df)
+
+        future = model.make_future_dataframe(
+            periods=args.periods, freq=args.freq, include_history=args.include_history
+        )
+        forecast = model.predict(future)
+
+        plot_path = self._save_forecast_plot(model, forecast, args.dataframe)
+
+        return ToolResult(
+            status="success",
+            result={
+                "forecast": forecast.to_dict(orient="records"),
+                "forecast_columns": list(forecast.columns),
+                "future_dataframe": future.to_dict(orient="records"),
+                "figure_path": str(plot_path),
+                "figure_url": self.to_static_url(plot_path),
+            },
+            metadata={
+                "dataframe": args.dataframe,
+                "ds_column": args.ds_column,
+                "y_column": args.y_column,
+                "periods": args.periods,
+                "freq": args.freq,
+            },
+        )
+
+    def _resolve_dataframe(self, name: str) -> pd.DataFrame:
+        if name in self.dataframes:
+            return self.dataframes[name]
+
+        # Accept alias names (df1, df2, etc.)
+        for df_name, alias in self.alias_map.items():
+            if name == alias and df_name in self.dataframes:
+                return self.dataframes[df_name]
+
+        raise ValueError(
+            f"DataFrame '{name}' not found. Available: {list(self.dataframes.keys())}"
+        )
+
+    def _prepare_dataframe(
+        self, dataframe: pd.DataFrame, args: ProphetForecastArgs
+    ) -> pd.DataFrame:
+        if args.ds_column not in dataframe.columns:
+            raise ValueError(
+                f"Date column '{args.ds_column}' not found. Columns: {list(dataframe.columns)}"
+            )
+        if args.y_column not in dataframe.columns:
+            raise ValueError(
+                f"Target column '{args.y_column}' not found. Columns: {list(dataframe.columns)}"
+            )
+
+        df = dataframe.copy()
+        df[args.ds_column] = pd.to_datetime(df[args.ds_column], errors="coerce")
+        df[args.y_column] = pd.to_numeric(df[args.y_column], errors="coerce")
+
+        df = df.dropna(subset=[args.ds_column, args.y_column])
+        if df.empty:
+            raise ValueError("No valid rows available after cleaning for Prophet forecast")
+
+        df = df[[args.ds_column, args.y_column]].rename(
+            columns={args.ds_column: "ds", args.y_column: "y"}
+        )
+        return df
+
+    def _save_forecast_plot(
+        self, model: Prophet, forecast: pd.DataFrame, dataframe_name: str
+    ) -> Path:
+        plot_dir = self.output_dir
+        plot_dir.mkdir(parents=True, exist_ok=True)
+
+        fig = model.plot(forecast)
+        filename = f"prophet_forecast_{dataframe_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+        file_path = plot_dir / filename
+        fig.savefig(file_path, bbox_inches="tight")
+        return file_path