diff --git a/plotsense.egg-info/PKG-INFO b/plotsense.egg-info/PKG-INFO index d9dde76..23b086a 100644 --- a/plotsense.egg-info/PKG-INFO +++ b/plotsense.egg-info/PKG-INFO @@ -1,10 +1,11 @@ Metadata-Version: 2.4 Name: plotsense -Version: 0.1.0 +Version: 0.1.3 Summary: An intelligent plotting package with suggestions and explanations -Author-email: Christian Chimezie -License: MIT -Project-URL: Homepage, https://github.com/christianchimezie/PlotSenseAI +Home-page: https://github.com/christianchimezie/PlotSenseAI +Author: Christian Chimezie, Toluwaleke Ogidan, Grace Farayola, Amaka Iduwe, Nelson Ogbeide, Onyekachukwu Ojumah, Olamilekan Ajao +Author-email: chimeziechristiancc@gmail.com, gbemilekeogidan@gmail.com, gracefarayola@gmail.com, nwaamaka_iduwe@yahoo.com, Ogbeide331@gmail.com, Onyekaojumah22@gmail.com, olamilekan011@gmail.com +License: Apache License 2.0 Project-URL: Documentation, https://github.com/christianchimezie/PlotSenseAI/blob/main/README.md Classifier: Development Status :: 3 - Alpha Classifier: Intended Audience :: Science/Research @@ -21,11 +22,28 @@ Requires-Python: >=3.8 Description-Content-Type: text/markdown License-File: LICENCE License-File: NOTICE -Requires-Dist: matplotlib>=3.0 +Requires-Dist: matplotlib>=3.8.0 Requires-Dist: seaborn>=0.11 Requires-Dist: pandas>=1.0 Requires-Dist: numpy>=1.18 +Requires-Dist: python-dotenv +Requires-Dist: groq +Requires-Dist: anthropic +Requires-Dist: openai +Requires-Dist: google-genai +Requires-Dist: requests +Dynamic: author +Dynamic: author-email +Dynamic: classifier +Dynamic: description +Dynamic: description-content-type +Dynamic: home-page +Dynamic: license Dynamic: license-file +Dynamic: project-url +Dynamic: requires-dist +Dynamic: requires-python +Dynamic: summary # 🌟 PlotSense: AI-Powered Data Visualization Assistant @@ -51,7 +69,7 @@ pip install plotsense ```bash import plotsense as ps -from plotsense import recommender, generate_plot, explainer, +from plotsense import recommender, plotgen, explainer ``` ### πŸ” Authenticate with Groq API: Get your free API key from Groq Cloud https://console.groq.com/home @@ -82,7 +100,7 @@ print(suggestions) ``` ### πŸ“Š Sample Output: -![alt text](suggestions_table.png) +![alt text](image.png) πŸŽ›οΈ Want more suggestions? @@ -90,7 +108,30 @@ print(suggestions) suggestions = ps.recommender(df, n=10) ``` -### 🧾 2. AI-Powered Plot Explanation +### πŸ“ˆ 2. One-Click Plot Generation +Generate recommended charts instantly: + +```bash +plot1 = ps.plotgen(df, suggestions.iloc[0]) # This will plot a bar chart with variables 'survived', 'pclass' +plot2 = ps.plotgen(df, suggestions.iloc[1]) # This will plot a bar chart with variables 'survived', 'sex' +plot3 = ps.plotgen(df, suggestions.iloc[2]) # This will plot a histogram with variable 'age' +``` +πŸŽ›οΈ Want more control? + +``` bash +plot1 = ps.plotgen(df, suggestions.iloc[0], x='pclass', y='survived') +``` +Supported Plots +- scatter +- bar +- barh +- histogram +- boxplot +- violinplot +- pie +- hexbin + +### 🧾 3. AI-Powered Plot Explanation Turn your visualizations into stories with natural language insights: ``` bash @@ -103,7 +144,7 @@ print(explanation) - Custom Prompts: You can provide your own prompt to guide the explanation ``` bash -explanation = refine_plot_explanation( +explanation = explainer( fig, prompt="Explain the key trends in this sales data visualization" ) @@ -111,7 +152,14 @@ explanation = refine_plot_explanation( - Multiple Refinement Iterations: Increase the number of refinement cycles for more polished explanations: ```bash -explanation = refine_plot_explanation(fig, iterations=3) # Default is 2 +explanation = explainer(fig, max_iterations=3) # Default is 2 +``` + +## πŸ”„ Combined Workflow: Suggest β†’ Plot β†’ Explain +``` bash +suggestions = ps.recommender(df) +plot = ps.plotgen(df, suggestions.iloc[0]) +insight = ps.explainer(plot) ``` ## 🀝 Contributing @@ -131,13 +179,15 @@ We welcome contributions! - More model integrations - Automated insight highlighting - Jupyter widget support +- Features/target analysis +- More supported plots ### πŸ“₯ Install or Update ``` bash pip install --upgrade plotsense # Get the latest features! ``` ## πŸ›‘ License -MIT License (Open Source) +Apache License 2.0 ## πŸ” API & Privacy Notes - Your API key is securely held in memory for your current Python session. @@ -146,3 +196,12 @@ MIT License (Open Source) Let your data speakβ€”with clarity, power, and PlotSense. πŸ“Šβœ¨ + +## Your Feedback +[Feedback Form](https://forms.gle/QEjipzHiMagpAQU99) + + + + + + diff --git a/plotsense.egg-info/SOURCES.txt b/plotsense.egg-info/SOURCES.txt index a251ad9..0edb61d 100644 --- a/plotsense.egg-info/SOURCES.txt +++ b/plotsense.egg-info/SOURCES.txt @@ -1,3 +1,5 @@ +LICENCE +NOTICE README.md pyproject.toml setup.py @@ -10,6 +12,24 @@ plotsense.egg-info/top_level.txt plotsense/explanations/__init__.py plotsense/explanations/explanations.py plotsense/plot_generator/__init__.py +plotsense/plot_generator/base_generator.py +plotsense/plot_generator/basic_generator.py plotsense/plot_generator/generator.py +plotsense/plot_generator/helpers.py +plotsense/plot_generator/registry.py +plotsense/plot_generator/smart_generator.py +plotsense/plot_generator/plots/__init__.py +plotsense/visual_suggestion/__init__.py plotsense/visual_suggestion/suggestions.py -plotsense/visual_suggestion/__init__.py \ No newline at end of file +plotsense/visual_suggestion/recommender/__init__.py +plotsense/visual_suggestion/recommender/dataframe_analyzer.py +plotsense/visual_suggestion/recommender/ensemble_scorer.py +plotsense/visual_suggestion/recommender/prompt_builder.py +plotsense/visual_suggestion/recommender/response_parser.py +plotsense/visual_suggestion/recommender/visualization_recommender.py +test/__init__.py +test/my_ptce_test.py +test/my_test.py +test/test_explanations.py +test/test_plotgen.py +test/test_suggestions.py \ No newline at end of file diff --git a/plotsense.egg-info/requires.txt b/plotsense.egg-info/requires.txt index 6fbc6f2..a329864 100644 --- a/plotsense.egg-info/requires.txt +++ b/plotsense.egg-info/requires.txt @@ -1,4 +1,10 @@ -matplotlib>=3.0 +matplotlib>=3.8.0 seaborn>=0.11 pandas>=1.0 numpy>=1.18 +python-dotenv +groq +anthropic +openai +google-genai +requests diff --git a/plotsense.egg-info/top_level.txt b/plotsense.egg-info/top_level.txt index 100c7e8..f5522c3 100644 --- a/plotsense.egg-info/top_level.txt +++ b/plotsense.egg-info/top_level.txt @@ -1 +1,2 @@ plotsense +test diff --git a/plotsense/__init__.py b/plotsense/__init__.py index 301ab30..9d8b1dd 100644 --- a/plotsense/__init__.py +++ b/plotsense/__init__.py @@ -1,3 +1,4 @@ from plotsense.visual_suggestion.suggestions import recommender, VisualizationRecommender -from plotsense.explanations.explanations import explainer,PlotExplainer -from plotsense.plot_generator.generator import plotgen, PlotGenerator \ No newline at end of file +from plotsense.explanations.explanations import explainer, PlotExplainer +from plotsense.plot_generator.generator import plotgen, PlotGenerator + diff --git a/plotsense/core/ai_interface.py b/plotsense/core/ai_interface.py new file mode 100644 index 0000000..eb17f22 --- /dev/null +++ b/plotsense/core/ai_interface.py @@ -0,0 +1,423 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +import warnings +from typing import Dict, List, Optional, Tuple + +from plotsense.core.enums.strategy import StrategyName +from plotsense.core.strategies.round_robin import RoundRobinStrategy +from plotsense.core.strategies.cost_optimized import CostOptimizedStrategy +from plotsense.core.strategies.performance_optimized import PerformanceOptimizedStrategy +from plotsense.core.strategies.fallback_chain import FallbackChainStrategy + + +class AIModelInterface: + """ + Handles all low-level interactions with LLM providers. + Acts as a bridge between PlotExplainer (or any client) + and ProviderManager. + """ + + def __init__(self, provider_manager, timeout: int = 30): + self.manager = provider_manager + self.timeout = timeout + + def _init_strategy( + self, strategy_name: StrategyName, + available_models: List[Tuple[str, str]] + ): + try: + strategy_name = StrategyName(strategy_name) + except ValueError: + raise ValueError(f"Invalid strategy name: {strategy_name}") + + if strategy_name == StrategyName.ROUND_ROBIN: + return RoundRobinStrategy(available_models) + elif strategy_name == StrategyName.COST_OPTIMIZED: + return CostOptimizedStrategy(available_models, self.manager) + elif strategy_name == StrategyName.PERFORMANCE: + return PerformanceOptimizedStrategy(available_models, self.manager) + elif strategy_name == StrategyName.FALLBACK_CHAIN: + return FallbackChainStrategy(available_models) + + def query_all_models( + self, + prompt: str, + debug: bool = False, + base64_image: Optional[str] = None, + custom_parameters: Optional[Dict] = None, + strategy: StrategyName = StrategyName.ROUND_ROBIN, + max_workers: int = 6, + ) -> Dict[str, str]: + """ + Query all available models (across all providers) in parallel. + Uses ThreadPoolExecutor for concurrency. + Returns a mapping of "provider:model" -> response_text. + + Notes: + - Keeps strategy initialization for compatibility (strategy instance + can be used later to reorder or filter models). + - Each model is queried independently; failures don't stop the rest. + """ + # Get available models as list of tuples: [(provider, model_name), ...] + all_models = self.manager.list_all_models() + self.available_models = [ + (provider, model) + for provider, models in all_models.items() + for model in models + ] + if not self.available_models: + raise ValueError("No available models found from provider manager.") + + # Initialize strategy instance (keeps previous behavior) + strategy_instance = self._init_strategy( + strategy, self.available_models + ) + + results: Dict[str, str] = {} + + # --- 1️⃣ Let strategy select or order models --- + # Most strategies (RoundRobin, CostOptimized, etc.) will implement a method + # like `.select_models(n: int)` or `.get_next_batch()`. + # If not, we simply use all available models. + try: + # Example interface: select_models returns a prioritized list + models_to_query = strategy_instance.select_model(len(self.available_models)) + except AttributeError: + # Fallback: strategy not yet implementing selection + models_to_query = self.available_models + + if not models_to_query: + raise ValueError("Strategy did not return any models to query.") + + if debug: + print(f"\n[DEBUG] Strategy '{strategy_instance.__class__.__name__}' selected models:") + for prov, mod in models_to_query: + print(f" - {prov}:{mod}") + + def _query_one(provider: str, model_name: str): + model_key = f"{provider}:{model_name}" + try: + resp = self.query_model( + provider=provider, + model=model_name, + prompt=prompt, + base64_image=base64_image, + custom_parameters=custom_parameters, + ) + return model_key, resp + except Exception as e: + warnings.warn(f"[AIModelInterface] Query failed for {model_key} -> {e}") + return model_key, f"Error: {e}" + + # FallbackChainStrategy -> sequential queries until one succeeds + if isinstance(strategy_instance, FallbackChainStrategy): + for provider, model_name in models_to_query: + key, resp = _query_one(provider, model_name) + results[key] = resp + if not resp.lower().startswith("error"): + # Stop at first success (fallback semantics) + break + else: + # Run queries concurrently + with ThreadPoolExecutor(max_workers=max_workers) as executor: + future_to_key = { + executor.submit(_query_one, provider, model_name): (provider, model_name) + for provider, model_name in self.available_models + } + + for future in as_completed(future_to_key): + key, resp = future.result() + results[key] = resp + + return results + + def query_model( + self, + provider: str, + model: str, + prompt: str, + base64_image: Optional[str] = None, + custom_parameters: Optional[Dict] = None + ) -> str: + """ + Query a model via the provider manager. + Handles provider-specific formatting and error management. + """ + if provider not in self.manager.providers: + raise ValueError(f"Unknown provider: {provider}") + + try: + # Build messages depending on provider/model + messages = self._build_messages( + provider, model, prompt, base64_image + ) + generation_params = {"temperature": 0.7, "max_tokens": 1000, **(custom_parameters or {})} + + provider_lower = provider.lower() + # model_lower = model.lower() + + # -------------------- OPENAI (Chat + Response) -------------------- + if "openai" in provider_lower: + # if "gpt" in model_lower or "chat" in model_lower: + if "chat" in provider_lower: + # Chat-based models (GPT-4, GPT-3.5, etc.) + return self.manager.query( + provider, + model=model, + messages=messages, + prompt=prompt, + **generation_params, + ) + elif "response" in provider_lower: + # Response-based models (completion endpoints) + return self.manager.query( + provider, + model=model, + prompt=prompt, + **generation_params, + ) + + # -------------------- AZURE OPENAI -------------------- + elif "azure" in provider_lower: + # Azure follows OpenAI's API style but requires deployment-specific name + return self.manager.query( + provider, + model=model, + messages=messages, + prompt=prompt, + **generation_params, + ) + + # -------------------- GROQ -------------------- + elif "groq" in provider_lower: + # Typically text-only Llama-style models + return self.manager.query( + provider, + model=model, + messages=messages, + prompt=prompt, + **generation_params, + ) + + # -------------------- ANTHROPIC -------------------- + elif "anthropic" in provider_lower: + # Claude models (text + multimodal optional) + return self.manager.query( + provider, + model=model, + messages=messages, + prompt=prompt, + **generation_params, + ) + + # -------------------- GEMINI -------------------- + elif "gemini" in provider_lower: + # Supports text + images + return self.manager.query( + provider, + model=model, + messages=messages, + prompt=prompt, + image=base64_image, + **generation_params, + ) + + # -------------------- OLLAMA -------------------- + elif "ollama" in provider_lower: + # Local models; prompt only, may support images if model allows + return self.manager.query( + provider, + model=model, + prompt=prompt, + image=base64_image, + **generation_params, + ) + + # -------------------- DEFAULT / UNKNOWN -------------------- + else: + print(f"[AIModelInterface] Warning: Using default query handling for {provider}:{model}") + # Fallback for new or custom providers + return self.manager.query( + provider, + model=model, + messages=messages, + prompt=prompt, + **generation_params, + ) + + except Exception as e: + warnings.warn(f"[AIModelInterface] Querying error for {provider}:{model} -> {str(e)}") + return f"Error: {e}" + finally: + return f"Error: No valid query handler found for provider '{provider}'." + + def get_model_weights(self) -> Dict[str, float]: + """ + Return model weighting for ensemble scoring. + + Weighting strategy (default heuristics): + - OpenAI GPT-4 variants -> higher weight (2.0) + - Anthropic Claude family -> high weight (1.8) + - Google Gemini -> high weight (1.6) + - Azure (OpenAI in Azure) -> treated similar to openai (1.8 for gpt-4) + - Groq (Llama variants) -> moderate weight (1.2) + - Ollama / local models -> lower/moderate weight (1.0) + - Other / unknown -> base weight (1.0) + + Returns: + dict of "provider:model" -> normalized_weight + """ + all_models = self.manager.list_all_models() + self.available_models = [ + (provider, model) + for provider, models in all_models.items() + for model in models + ] + + raw_weights: Dict[str, float] = {} + + for provider, model_name in self.available_models: + key = f"{provider}:{model_name}" + lname = model_name.lower() + lprov = provider.lower() + + # Base preference by model name + if "gpt-4" in lname or "gpt4" in lname or "gpt-4o" in lname: + base = 2.0 + elif "claude" in lname: + base = 1.8 + elif "gemini" in lname or "gemini-pro" in lname: + base = 1.6 + elif "llama" in lname or "groq" in lprov: + # groq's Llama-based models - decent but not highest + base = 1.2 + elif "azure" in lprov: + # Azure OpenAI often runs OpenAI models; favor if contains gpt-4 + base = 1.8 if "gpt-4" in lname or "gpt4" in lname else 1.1 + elif "ollama" in lprov: + base = 1.0 + else: + base = 1.0 + + # Provider-level adjustments (optional) + if lprov == "anthropic": + base *= 1.0 # already accounted by 'claude' checks + if lprov == "openai": + base *= 1.0 + if lprov == "groq": + base *= 1.0 + if lprov == "azure": + base *= 1.0 + + raw_weights[key] = base + + # Normalize to sum to 1 + total = sum(raw_weights.values()) or 1.0 + normalized = {k: (v / total) for k, v in raw_weights.items()} + return normalized + + def _build_messages( + self, provider: str, model: str, prompt: str, + base64_image: Optional[str] = None + ): + """ + Build messages dynamically based on provider capabilities. + Supports multimodal input where possible (OpenAI GPT-4o, Gemini, Anthropic, etc.). + Falls back to text-only prompt for providers without image support. + """ + provider_lower = provider.lower() + model_lower = model.lower() + + # --- 1️⃣ OpenAI / Azure (GPT-4, GPT-4o, GPT-3.5 etc.) --- + if provider_lower in {"openai", "azure"}: + if base64_image and any(tag in model_lower for tag in ["gpt-4o", "gpt-4-turbo", "gpt-4-vision"]): + # Chat message with multimodal support + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}, + ], + } + ] + else: + # Standard chat completion format + return [ + {"role": "system", "content": "You are a helpful data visualization assistant."}, + {"role": "user", "content": prompt}, + ] + + # --- 2️⃣ Anthropic (Claude) --- + elif provider_lower == "anthropic": + # Claude supports multimodal via text + image blocks in messages + if base64_image: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64_image}}, + ], + } + ] + else: + return [ + {"role": "user", "content": prompt} + ] + + # --- 3️⃣ Gemini (Google) --- + elif provider_lower == "gemini": + # Gemini API supports multimodal via a combined structure + if base64_image: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image", "data": base64_image, "mime_type": "image/jpeg"}, + ], + } + ] + else: + return [ + {"role": "user", "content": prompt} + ] + + # --- 4️⃣ Groq (LLaMA / Mistral etc. – text-only) --- + elif provider_lower == "groq": + return [ + {"role": "user", "content": prompt} + ] + + # --- 5️⃣ Ollama (local models; may support image, but prompt-based) --- + elif provider_lower == "ollama": + if base64_image: + # Send inline text prompt mentioning image context + return [ + { + "role": "user", + "content": f"{prompt}\n\n[Image attached as base64 input]" + } + ] + else: + return [ + {"role": "user", "content": prompt} + ] + + # --- 6️⃣ Default / Unknown provider fallback --- + else: + if base64_image: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}, + ], + } + ] + else: + return [ + {"role": "user", "content": prompt} + ] + diff --git a/plotsense/core/enums/strategy.py b/plotsense/core/enums/strategy.py new file mode 100644 index 0000000..61dbecc --- /dev/null +++ b/plotsense/core/enums/strategy.py @@ -0,0 +1,8 @@ +from enum import Enum + +class StrategyName(str, Enum): + ROUND_ROBIN = "round_robin" + COST_OPTIMIZED = "cost_optimized" + PERFORMANCE = "performance" + FALLBACK_CHAIN = "fallback" + diff --git a/plotsense/core/providers/anthropic.py b/plotsense/core/providers/anthropic.py new file mode 100644 index 0000000..a24e290 --- /dev/null +++ b/plotsense/core/providers/anthropic.py @@ -0,0 +1,75 @@ +from typing import List +from anthropic import Anthropic +from .base import LLMProvider + + +class AnthropicProvider(LLMProvider): + """Provider integration for Anthropic's Claude models.""" + + LINK = "πŸ‘‰ https://console.anthropic.com/account/keys πŸ‘ˆ" + + def __init__(self, api_key: str): + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize Anthropic client if not already created.""" + if not self.client: + self.client = Anthropic(api_key=self.api_key) + + def query(self, prompt: str, model: str, **kwargs) -> str: + """Send a message to Anthropic Claude model and return its response text.""" + if not self.client: + raise ValueError( + "Anthropic client not initialized. Call validate_key() first." + ) + + messages = kwargs.pop("messages", None) + if not messages and prompt: + # Default to a single user message + messages = [{"role": "user", "content": prompt}] + elif not messages and not prompt: + raise ValueError("Either 'prompt' or 'messages' must be provided.") + + try: + response = self.client.messages.create( + model=model, + messages=[{"role": "user", "content": prompt}], + **kwargs, + ) + return response.content[0].text if response and response.content else "" + except Exception as e: + raise RuntimeError(f"Anthropic query failed: {e}") + + def list_models(self) -> List[str]: + """ + Return a list of supported Anthropic models. + This list can be expanded as new Claude versions are released. + """ + return [ + "claude-3-5-sonnet-20241022", + "claude-3-opus-20240229", + "claude-3-haiku-20240307", + ] + + def validate_key(self) -> bool: + """ + Validate the provided API key by performing a lightweight test. + Returns True if successful, False otherwise. + """ + try: + self._init_client() + if not self.client: + raise ValueError( + "Anthropic client not initialized. Call validate_key() first." + ) + # Perform a trivial, cheap call to verify authentication + self.client.messages.create( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "ping"}], + max_tokens=5, + ) + return True + except Exception: + return False + diff --git a/plotsense/core/providers/azure_openai.py b/plotsense/core/providers/azure_openai.py new file mode 100644 index 0000000..4a81200 --- /dev/null +++ b/plotsense/core/providers/azure_openai.py @@ -0,0 +1,98 @@ +from typing import List +from openai import OpenAI +# AzureOpenAI, +from openai.types.chat import ChatCompletionUserMessageParam +from .base import LLMProvider + + +class AzureOpenAIProvider(LLMProvider): + """Provider integration for Azure-hosted OpenAI models.""" + + LINK = "πŸ‘‰ https://portal.azure.com/#view/Microsoft_Azure_ProjectOxford/CognitiveServicesHub/feature/OpenAI πŸ‘ˆ" + + def __init__( + self, api_key: str, + endpoint: str = "https://models.github.ai/inference", + api_version: str = "2024-02-15-preview" + ): + """ + Args: + api_key: Azure OpenAI API key + endpoint: Full Azure endpoint (e.g. https://.openai.azure.com/) + api_version: Azure OpenAI API version + """ + self.api_key = api_key + self.endpoint = endpoint + # self.api_version = api_version + self.client = None + + def _init_client(self): + """Initialize the Azure OpenAI client.""" + if not self.endpoint: + raise ValueError("Azure OpenAI endpoint not provided.") + if not self.client: + self.client = OpenAI( + api_key=self.api_key, + # api_version=self.api_version, + # azure_endpoint=self.endpoint, + base_url=self.endpoint, + ) + + def query(self, prompt: str, model: str, **kwargs) -> str: + """ + Send a prompt to Azure OpenAI Chat Completion API. + """ + self._init_client() + if not self.client: + raise ValueError("AzureOpenAI client not initialized. Call validate_key() first.") + + # Ensure messages format exists in kwargs + messages: list[ChatCompletionUserMessageParam] = kwargs.pop("messages", None) + if not messages and prompt: + messages = [{"role": "user", "content": prompt}] + elif not messages and not prompt: + raise ValueError("Either 'prompt' or 'messages' must be provided.") + + try: + if "max_tokens" in kwargs: + kwargs["max_output_tokens"] = kwargs.pop("max_tokens") + response = self.client.chat.completions.create( + model=model, + messages=messages, + **kwargs + ) + return response.choices[0].message.content + except Exception as e: + raise RuntimeError(f"Azure OpenAI query failed: {e}") + + def list_models(self) -> List[str]: + """ + Return a suggested list of Azure OpenAI deployable model names. + (These must match your deployment names in Azure.) + """ + return [ + "openai/gpt-5", + # "gpt-4o", + # "gpt-4-turbo", + # "gpt-35-turbo", + # "gpt-4", + ] + + def validate_key(self) -> bool: + """ + Attempt a lightweight ping to validate Azure OpenAI credentials. + """ + try: + self._init_client() + if not self.client: + raise ValueError("AzureOpenAI client not initialized. Call validate_key() first.") + response = self.client.chat.completions.create( + model="openai/gpt-5", + messages=[{"role": "user", "content": "ping"}], + max_completion_tokens=5 + ) + return bool(response) + except Exception as e: + print(f"⚠️ Azure OpenAI API key validation failed: {e}") + return False + diff --git a/plotsense/core/providers/base.py b/plotsense/core/providers/base.py new file mode 100644 index 0000000..e1f8d94 --- /dev/null +++ b/plotsense/core/providers/base.py @@ -0,0 +1,25 @@ +from abc import ABC, abstractmethod +from typing import List + +class LLMProvider(ABC): + """Abstract base class for LLM providers.""" + + LINK: str + + @abstractmethod + def __init__(self, api_key: str): + """Initialize the provider with an API key.""" + pass + + @abstractmethod + def query(self, prompt: str, model: str, **kwargs) -> str: + pass + + @abstractmethod + def list_models(self) -> List[str]: + pass + + @abstractmethod + def validate_key(self) -> bool: + pass + diff --git a/plotsense/core/providers/gemini.py b/plotsense/core/providers/gemini.py new file mode 100644 index 0000000..5d3d791 --- /dev/null +++ b/plotsense/core/providers/gemini.py @@ -0,0 +1,92 @@ +from typing import List, Optional +from google import genai +from google.genai.types import GenerateContentConfig +from .base import LLMProvider + + +class GeminiProvider(LLMProvider): + """Provider integration for Google's Gemini models (v2 SDK).""" + + LINK = "πŸ‘‰ https://aistudio.google.com/app/apikey πŸ‘ˆ" + + def __init__(self, api_key: str): + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize Anthropic client if not already created.""" + if not self.client: + self.client = genai.Client(api_key=self.api_key) + + def query( + self, + prompt: str, + model: str, + base64_image: Optional[str] = None, + **kwargs, + ) -> str: + """ + Send a text (or multimodal) prompt to Gemini and return the response text. + Supports text-only and text+image queries. + + Uses the new google-genai v2 API. + """ + try: + # Build input depending on image presence + if base64_image: + # Multimodal: send both text and image + contents = [ + {"text": prompt}, + { + "inline_data": { + "mime_type": "image/jpeg", + "data": base64_image, + } + }, + ] + else: + # Text-only + contents = prompt + + self._init_client() + if not self.client: + raise ValueError("Gemini client initialization failed.") + + response = self.client.models.generate_content( + model=model, + contents=contents, + **kwargs, + ) + + # Return clean text or empty string if missing + return getattr(response, "text", "") or "" + + except Exception as e: + raise RuntimeError(f"Gemini query failed: {e}") + + def list_models(self) -> List[str]: + """ + Return a curated list of Gemini models. + """ + return [ + "gemini-2.5-flash", + "gemini-2.0-pro", + "gemini-1.5-flash", + ] + + def validate_key(self) -> bool: + """ + Validate the provided Gemini API key by attempting a trivial generation. + """ + try: + self._init_client() + if not self.client: + raise ValueError("Gemini client initialization failed.") + response = self.client.models.generate_content( + model="gemini-2.5-flash", + contents="ping", + config=GenerateContentConfig(max_output_tokens=5), + ) + return bool(response.text) + except Exception: + return False diff --git a/plotsense/core/providers/groq.py b/plotsense/core/providers/groq.py new file mode 100644 index 0000000..abbb412 --- /dev/null +++ b/plotsense/core/providers/groq.py @@ -0,0 +1,74 @@ +from typing import List +from groq import Groq +from groq.types.chat import ChatCompletionUserMessageParam +from .base import LLMProvider + + +class GroqProvider(LLMProvider): + """Provider integration for Groq's fast inference API.""" + + LINK = "πŸ‘‰ https://console.groq.com/keys πŸ‘ˆ" + + def __init__(self, api_key: str): + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize Groq client if not already created.""" + if not self.client: + self.client = Groq(api_key=self.api_key) + + def query( + self, + prompt: str, + model: str, + **kwargs, + ) -> str: + """ + Send a text prompt to Groq (Llama models) and return the response text. + Supports OpenAI-style chat completions. + """ + self._init_client() + if not self.client: + raise ValueError("Groq client not initialized. Call validate_key() first.") + + # Build messages dynamically (fallback if only prompt is given) + messages: list[ChatCompletionUserMessageParam] = kwargs.pop("messages", None) + if not messages and prompt: + messages = [{"role": "user", "content": prompt}] + elif not messages and not prompt: + raise ValueError("Either 'prompt' or 'messages' must be provided.") + + try: + response = self.client.chat.completions.create( + model=model, + messages=messages, + **kwargs, + ) + return response.choices[0].message.content + except Exception as e: + raise RuntimeError(f"Groq query failed: {e}") + + def list_models(self) -> List[str]: + """Return a curated list of supported Groq models.""" + return [ + "llama-3.1-8b-instant", + "llama-3.3-70b-versatile", + ] + + def validate_key(self) -> bool: + """ + Validate the provided Groq API key by making a lightweight request. + """ + try: + self._init_client() + if not self.client: + raise ValueError("Groq client not initialized.") + response = self.client.chat.completions.create( + model="llama-3.1-8b-instant", + messages=[{"role": "user", "content": "ping"}], + max_tokens=5, + ) + return bool(response.choices[0].message.content) + except Exception: + return False diff --git a/plotsense/core/providers/groq_openai.py b/plotsense/core/providers/groq_openai.py new file mode 100644 index 0000000..ca0f4f7 --- /dev/null +++ b/plotsense/core/providers/groq_openai.py @@ -0,0 +1,73 @@ +from typing import List +from openai.types.chat import ChatCompletionUserMessageParam +from openai import OpenAI +from .base import LLMProvider + + +class GroqProvider(LLMProvider): + """Provider for Groq models using the unified OpenAI SDK interface.""" + + LINK = "πŸ‘‰ https://console.groq.com/keys πŸ‘ˆ" + + def __init__(self, api_key: str): + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize Groq client via OpenAI-compatible endpoint.""" + if not self.client: + self.client = OpenAI( + api_key=self.api_key, + base_url="https://api.groq.com/openai/v1" # Key difference + ) + + def query(self, prompt: str, model: str, **kwargs) -> str: + """ + Send a chat completion query to Groq via OpenAI SDK. + """ + self._init_client() + if not self.client: + raise ValueError("Groq client not initialized. Call validate_key() first.") + + # Ensure messages are present + messages: list[ChatCompletionUserMessageParam] = kwargs.pop( + "messages", None + ) + if not messages and prompt: + messages = [{"role": "user", "content": prompt}] + elif not messages and not prompt: + raise ValueError("Either 'prompt' or 'messages' must be provided.") + + try: + response = self.client.chat.completions.create( + model=model, + messages=messages, + **kwargs + ) + return response.choices[0].message.content + except Exception as e: + raise RuntimeError(f"Groq query failed: {e}") + + def list_models(self) -> List[str]: + """ + Available Groq Llama models (you can update this dynamically later). + """ + return ["llama-3.1-8b-instant", "llama-3.3-70b-versatile"] + + def validate_key(self) -> bool: + """ + Simple ping to check API validity. + """ + try: + self._init_client() + if not self.client: + raise ValueError("Groq OpenAI client not initialized.") + response = self.client.chat.completions.create( + model="llama-3.1-8b-instant", + messages=[{"role": "user", "content": "ping"}], + max_tokens=5 + ) + return bool(response) + except Exception: + return False + diff --git a/plotsense/core/providers/ollama_openai.py b/plotsense/core/providers/ollama_openai.py new file mode 100644 index 0000000..185a045 --- /dev/null +++ b/plotsense/core/providers/ollama_openai.py @@ -0,0 +1,76 @@ +from typing import List +from openai import OpenAI +from openai.types.chat import ChatCompletionUserMessageParam +from .base import LLMProvider + + +class OllamaProvider(LLMProvider): + """ + Provider for Ollama models using the OpenAI-compatible API. + This allows querying a locally running Ollama instance. + """ + + LINK = "πŸ‘‰ https://ollama.ai/library πŸ‘ˆ" + + def __init__(self, api_key: str = ""): + # Ollama typically doesn't require an API key (local service) + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize the OpenAI-compatible client for local Ollama.""" + if not self.client: + # Default local Ollama endpoint + self.client = OpenAI( + base_url="http://localhost:11434/v1", # Ollama’s OpenAI-compatible API + api_key=self.api_key or "ollama", # Dummy key for OpenAI client compatibility + ) + + def query(self, prompt: str, model: str, **kwargs) -> str: + """ + Query the Ollama model using OpenAI-compatible endpoint. + """ + self._init_client() + if not self.client: + raise ValueError("Ollama client not initialized. Call validate_key() first.") + + messages: list[ChatCompletionUserMessageParam] = kwargs.pop( + "messages", None + ) + if not messages and prompt: + messages = [{"role": "user", "content": prompt}] + elif not messages and not prompt: + raise ValueError("Either 'prompt' or 'messages' must be provided.") + + try: + response = self.client.chat.completions.create( + model=model, + **kwargs + ) + return response.choices[0].message.content + except Exception as e: + raise RuntimeError(f"Ollama query failed: {e}") + + def list_models(self) -> List[str]: + """ + List of example models. In a real setup, this could query `ollama list`. + """ + return ["llama3", "mistral", "codellama", "phi3", "neural-chat"] + + def validate_key(self) -> bool: + """ + Validate connection to local Ollama instance. + """ + try: + self._init_client() + if not self.client: + raise ValueError("Ollama OpenAI client not initialized.") + response = self.client.chat.completions.create( + model="llama3", + messages=[{"role": "user", "content": "ping"}], + max_tokens=5 + ) + return bool(response) + except Exception: + return False + diff --git a/plotsense/core/providers/openai_chat.py b/plotsense/core/providers/openai_chat.py new file mode 100644 index 0000000..e4bcd39 --- /dev/null +++ b/plotsense/core/providers/openai_chat.py @@ -0,0 +1,76 @@ +from typing import List, Optional +from openai import OpenAI +from openai.types.chat import ChatCompletionUserMessageParam +from .base import LLMProvider + + +class OpenAIChatProvider(LLMProvider): + """Provider integration for OpenAI Chat models.""" + + LINK = "πŸ‘‰ https://platform.openai.com/api-keys πŸ‘ˆ" + + def __init__(self, api_key: str): + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize OpenAI client if not already created.""" + if not self.client: + self.client = OpenAI(api_key=self.api_key) + + def query( + self, + prompt: Optional[str], + model: str, + **kwargs, + ) -> str: + """ + Send a prompt or messages to OpenAI Chat Completion API. + Supports both chat-style input and single text prompts. + """ + self._init_client() + if not self.client: + raise ValueError("OpenAI client not initialized. Call validate_key() first.") + + # Handle either prompt or messages + messages: list[ChatCompletionUserMessageParam] = kwargs.pop("messages", None) + if not messages and prompt: + messages = [{"role": "user", "content": prompt}] + elif not messages and not prompt: + raise ValueError("Either 'prompt' or 'messages' must be provided.") + + try: + response = self.client.chat.completions.create( + model=model, + messages=messages, + **kwargs, + ) + return response.choices[0].message.content + except Exception as e: + raise RuntimeError(f"OpenAI chat query failed: {e}") + + def list_models(self) -> List[str]: + """Return a curated list of supported OpenAI chat models.""" + return [ + "gpt-4o-mini", + "gpt-4.1", + "gpt-4-turbo", + "gpt-4o", + ] + + def validate_key(self) -> bool: + """ + Validate the provided OpenAI API key by performing a lightweight test query. + """ + try: + self._init_client() + if not self.client: + raise ValueError("OpenAI client not initialized.") + response = self.client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "ping"}], + max_tokens=5, + ) + return bool(response.choices[0].message.content) + except Exception: + return False diff --git a/plotsense/core/providers/openai_response.py b/plotsense/core/providers/openai_response.py new file mode 100644 index 0000000..bd4ec25 --- /dev/null +++ b/plotsense/core/providers/openai_response.py @@ -0,0 +1,77 @@ +from typing import List, Optional +from openai import OpenAI +from .base import LLMProvider + + +class OpenAIResponseProvider(LLMProvider): + """Provider integration for OpenAI's Responses API.""" + + LINK = "πŸ‘‰ https://platform.openai.com/api-keys πŸ‘ˆ" + + def __init__(self, api_key: str): + self.api_key = api_key + self.client = None + + def _init_client(self): + """Initialize OpenAI client if not already created.""" + if not self.client: + self.client = OpenAI(api_key=self.api_key) + + def query( + self, + prompt: Optional[str], + model: str, + **kwargs, + ) -> str: + """ + Send a prompt to OpenAI's Responses API and return the generated text. + The Responses endpoint supports text, image, and JSON outputs. + """ + self._init_client() + if not self.client: + raise ValueError("OpenAI client not initialized. Call validate_key() first.") + + if not prompt: + raise ValueError("'prompt' must be provided for Responses API queries.") + + try: + if "max_tokens" in kwargs: + kwargs["max_output_tokens"] = kwargs.pop("max_tokens") + # The Responses API expects `input`, not `messages` + response = self.client.responses.create( + model=model, + input=prompt, + **kwargs, + ) + + # The unified field for plain text output is `output_text` + return getattr(response, "output_text", "") or "" + except Exception as e: + raise RuntimeError(f"OpenAI response query failed: {e}") + + def list_models(self) -> List[str]: + """Return a curated list of supported OpenAI Response models.""" + return [ + "gpt-4.1", + "gpt-4.1-mini", + "gpt-4o", + ] + + def validate_key(self) -> bool: + """ + Validate the provided OpenAI API key by performing a lightweight test query. + """ + try: + self._init_client() + if not self.client: + raise ValueError("OpenAI client not initialized.") + response = self.client.responses.create( + model="gpt-4.1-mini", + input="ping", + max_output_tokens=16, + ) + return bool(getattr(response, "output_text", "")) + except Exception as e: + print(f"OpenAI Responses API key validation failed: {e}") + return False + diff --git a/plotsense/core/providers/provider_manager.py b/plotsense/core/providers/provider_manager.py new file mode 100644 index 0000000..fc50e4b --- /dev/null +++ b/plotsense/core/providers/provider_manager.py @@ -0,0 +1,241 @@ +from typing import Dict, List, Optional, Type + +from plotsense.core.providers.anthropic import AnthropicProvider +from plotsense.core.providers.azure_openai import AzureOpenAIProvider +from plotsense.core.providers.base import LLMProvider +from plotsense.core.providers.gemini import GeminiProvider +from plotsense.core.providers.ollama_openai import OllamaProvider +from plotsense.core.providers.openai_chat import OpenAIChatProvider +from plotsense.core.utils import prompt_for_api_key +from .groq import GroqProvider +from .openai_response import OpenAIResponseProvider + + +class ProviderManager: + """Manages multiple LLM providers, their API keys, and interactions.""" + + SUPPORTED_PROVIDERS: Dict[str, Dict[str, Type[LLMProvider]]] = { + "groq": { + "default": GroqProvider, + }, + "openai": { + "chat": OpenAIChatProvider, + "response": OpenAIResponseProvider, + }, + "anthropic": { + "default": AnthropicProvider, + }, + "gemini": { + "default": GeminiProvider, + }, + "azure": { + "default": AzureOpenAIProvider, + }, + "ollama": { + "default": OllamaProvider, + }, + } + + def __init__( + self, api_keys: Dict[str, str], interactive: bool = True, + restrict_to: Optional[List[str]] = None + ): + self.api_keys = api_keys or {} + self.interactive = interactive + self.providers = {} + self.restrict_to = set(restrict_to) if restrict_to else None + + # Normalize restrict_to list + if restrict_to: + invalid = [p for p in restrict_to if p not in self.SUPPORTED_PROVIDERS] + if invalid: + raise ValueError( + f"Unsupported provider(s): {invalid}. " + f"Supported providers: {list(self.SUPPORTED_PROVIDERS.keys())}" + ) + self.restrict_to = set(restrict_to) + else: + self.restrict_to = None + + self._init_providers() + + def _init_providers(self): + """Initialize all registered providers and validate their API keys.""" + for vendor_name, variants in self.SUPPORTED_PROVIDERS.items(): + # Skip if restrict_to is provided and this vendor isn’t included + if self.restrict_to and vendor_name not in self.restrict_to: + continue + + for variant_name, provider_cls in variants.items(): + full_name = f"{vendor_name}_{variant_name}" + link = getattr(provider_cls, "LINK", f"https://{vendor_name}.com") + + api_key: Optional[str] = self.api_keys.get(vendor_name) + if not api_key: + # Try to prompt only if interactive and not restricted + api_key = prompt_for_api_key( + vendor_name, + link, + self.interactive, + skip_if_missing=bool(self.restrict_to), + ) + if not api_key: + # Skip this provider if key is still missing + print(f"⏩ Skipping {full_name.upper()} (no API key provided).") + continue + + self.api_keys[vendor_name] = api_key + + if not isinstance(api_key, str) or not api_key.strip(): + print(f"⚠️ Skipping {full_name.upper()} due to invalid API key format.") + continue + + provider = provider_cls(api_key=api_key) + + try: + if provider.validate_key(): + print(f"βœ… {full_name.upper()} API key validated successfully.") + self.providers[full_name] = provider + else: + print(f"❌ {full_name.upper()} API key invalid or unverified.") + except Exception as e: + print(f"⚠️ Error validating {full_name.upper()} API key: {e}") + + def get_provider(self, vendor_name: str, variant_name: str = ""): + """ + Get or initialize a provider (with optional variant) on demand. + + Args: + vendor_name: Name of the AI provider (e.g., "openai", "groq") + variant_name: Optional variant name (e.g., "chat", "completion") + + Returns: + Initialized provider instance + """ + # Compose a unique key for storage + full_name = f"{vendor_name}_{variant_name}" if variant_name else vendor_name + + if vendor_name not in self.SUPPORTED_PROVIDERS: + raise ValueError(f"Unknown provider: {vendor_name}") + + if full_name not in self.providers: + variants = self.SUPPORTED_PROVIDERS[vendor_name] + + # Determine class safely + provider_cls = None + if variant_name: + provider_cls = variants.get(variant_name) + if not provider_cls: + raise ValueError( + f"Unknown variant '{variant_name}' for provider '{vendor_name}'" + ) + else: + variant_name, provider_cls = next(iter(variants.items())) + full_name = f"{vendor_name}_{variant_name}" + + link = getattr(provider_cls, "LINK", f"https://{vendor_name}.com") + + api_key: Optional[str] = self.api_keys.get(vendor_name) + if not api_key: + api_key = prompt_for_api_key(vendor_name, link, self.interactive) + if not api_key: + raise ValueError(f"Missing API key for {vendor_name}") + self.api_keys[vendor_name] = api_key + + # if not isinstance(api_key, str): + # raise TypeError(f"API key for {vendor_name} must be a string") + + provider = provider_cls(api_key=api_key) + + try: + if provider.validate_key(): + print(f"βœ… {full_name.upper()} API key validated successfully.") + else: + print(f"❌ {full_name.upper()} API key invalid or unverified.") + except Exception as e: + print(f"⚠️ Error validating {full_name.upper()} API key: {e}") + + self.providers[full_name] = provider + + return self.providers[full_name] + + def list_all_models(self): + all_models = {} + for name, provider in self.providers.items(): + try: + all_models[name] = provider.list_models() + except Exception as e: + print(f"⚠️ Failed to list models for {name}: {str(e)}") + return all_models + + def query(self, provider_name: str, model: str, prompt: str, **kwargs): + """Query a specific provider with a prompt and model.""" + provider = self.providers.get(provider_name) + if not provider: + raise ValueError(f"Provider {provider_name} not initialized.") + return provider.query(prompt, model, **kwargs) + + def get_model_costs(self) -> Dict[str, float]: + """ + Return a global map of model names to approximate per-request cost multipliers. + This helps CostOptimizedStrategy prioritize cheaper models. + """ + # In a real system, this could come from provider-specific metadata + return { + # OpenAI + "gpt-4o-mini": 0.01, + "gpt-4o": 0.03, + "gpt-4-turbo": 0.025, + "gpt-3.5-turbo": 0.008, + # Groq (Llama) + "llama-3.1-8b-instant": 0.005, + "llama-3.3-70b-versatile": 0.02, + # Anthropic + "claude-3-haiku": 0.009, + "claude-3-sonnet": 0.02, + "claude-3-opus": 0.05, + # Gemini + "gemini-1.5-flash": 0.006, + "gemini-1.5-pro": 0.02, + # Azure (proxy to GPT costs) + "azure-gpt-4o-mini": 0.011, + "azure-gpt-4o": 0.031, + # Ollama (local = near-zero cost) + "llama3": 0.001, + "mistral": 0.002, + } + + def get_model_performance(self) -> Dict[str, float]: + """ + Return approximate relative performance scores for each model. + Higher means better performance (accuracy, reasoning ability, etc.). + """ + return { + # OpenAI + "gpt-4o": 10.0, + "gpt-4o-mini": 8.5, + "gpt-4-turbo": 9.5, + "gpt-3.5-turbo": 7.5, + + # Anthropic + "claude-3-opus": 9.8, + "claude-3-sonnet": 9.0, + "claude-3-haiku": 7.0, + + # Groq + "llama-3.3-70b-versatile": 8.8, + "llama-3.1-8b-instant": 6.5, + + # Gemini + "gemini-1.5-pro": 9.3, + "gemini-1.5-flash": 7.8, + + # Azure (maps to OpenAI) + "azure-gpt-4o": 9.8, + "azure-gpt-4o-mini": 8.3, + + # Ollama (local models) + "mistral": 6.0, + "llama3": 6.8, + } + diff --git a/plotsense/core/strategies/cost_optimized.py b/plotsense/core/strategies/cost_optimized.py new file mode 100644 index 0000000..dd3b7a5 --- /dev/null +++ b/plotsense/core/strategies/cost_optimized.py @@ -0,0 +1,33 @@ +from typing import Dict, List, Optional, Tuple +from plotsense.core.strategies.strategy import Strategy + + +class CostOptimizedStrategy(Strategy): + """Prioritize cheaper models first, fallback to pricier if needed.""" + + def __init__(self, provider_models: List[Tuple[str, str]], provider_manager): + super().__init__(provider_models) + + self.cost_map: Dict[str, float] = provider_manager.get_model_costs() + + # Sort models by ascending cost (lowest first) + self.model_list = sorted( + provider_models, + key=lambda p_m: self.cost_map.get(p_m[1], float("inf")) + ) + + def select_models(self, n: int) -> List[Tuple[str, str]]: + """ + Return the top `n` cheapest models. + """ + return self.model_list[:n] + + def select_model( + self, iteration: int, current_explanation: Optional[str] = None + ) -> Tuple[str, str]: + """Use the cheapest available model, escalate if iteration increases.""" + if not self.model_list: + raise ValueError("No models available in strategy.") + index = min(iteration, len(self.model_list) - 1) + return self.model_list[index] + diff --git a/plotsense/core/strategies/fallback_chain.py b/plotsense/core/strategies/fallback_chain.py new file mode 100644 index 0000000..2110e8d --- /dev/null +++ b/plotsense/core/strategies/fallback_chain.py @@ -0,0 +1,24 @@ +from typing import List, Optional, Tuple +from plotsense.core.strategies.strategy import Strategy + +class FallbackChainStrategy(Strategy): + """Try providers/models in fixed order until one succeeds.""" + + def __init__(self, provider_models: List[Tuple[str, str]]): + super().__init__(provider_models) + # Deterministic order; could later be made configurable + self.model_list = provider_models + self._last_success_index = 0 + + def select_model( + self, iteration: int, current_explanation: Optional[str] = None + ) -> Tuple[str, str]: + """If previous success exists, keep using it; otherwise go to next.""" + if not self.model_list: + raise ValueError("No models available in strategy.") + index = min(iteration, len(self.model_list) - 1) + return self.model_list[index] + + def report_success(self, index: int): + """Optionally record which model last succeeded.""" + self._last_success_index = index diff --git a/plotsense/core/strategies/performance_optimized.py b/plotsense/core/strategies/performance_optimized.py new file mode 100644 index 0000000..bff04bb --- /dev/null +++ b/plotsense/core/strategies/performance_optimized.py @@ -0,0 +1,40 @@ +from typing import Dict, List, Optional, Tuple +from plotsense.core.strategies.strategy import Strategy + +MODEL_PERFORMANCE_MAP = { + "gpt-4o": 10, + "gpt-4o-mini": 8, + "llama-3.3-70b-versatile": 9, + "llama-3.1-8b-instant": 6, +} + +class PerformanceOptimizedStrategy(Strategy): + """Prefer highest-performance models first.""" + + def __init__( + self, provider_models: List[Tuple[str, str]], provider_manager + ): + super().__init__(provider_models) + + # Get dynamic performance scores from ProviderManager + self.performance_map: Dict[str, float] = provider_manager.get_model_performance() + + # Sort models descending by performance score + self.model_list = sorted( + provider_models, + key=lambda p_m: self.performance_map.get(p_m[1], 0), + reverse=True, + ) + + def select_models(self, n: int) -> List[Tuple[str, str]]: + """Return the top `n` highest-performing models.""" + return self.model_list[:n] + + def select_model( + self, iteration: int, current_explanation: Optional[str] = None + ) -> Tuple[str, str]: + """Start from best model; fallback to lower-tier ones if needed.""" + if not self.model_list: + raise ValueError("No models available in strategy.") + index = min(iteration, len(self.model_list) - 1) + return self.model_list[index] diff --git a/plotsense/core/strategies/round_robin.py b/plotsense/core/strategies/round_robin.py new file mode 100644 index 0000000..45a3c89 --- /dev/null +++ b/plotsense/core/strategies/round_robin.py @@ -0,0 +1,19 @@ +from typing import List, Optional, Tuple +from plotsense.core.strategies.strategy import Strategy + +class RoundRobinStrategy(Strategy): + """Cycle through all models evenly.""" + + def __init__(self, provider_models: List[Tuple[str, str]]): + super().__init__(provider_models) + self.model_list = provider_models + self._last_index = -1 + + def select_model( + self, iteration: int, current_explanation: Optional[str] = None + ) -> Tuple[str, str]: + if not self.model_list: + raise ValueError("No models available in strategy.") + # Pick model based directly on iteration count + index = iteration % len(self.model_list) + return self.model_list[index] diff --git a/plotsense/core/strategies/strategy.py b/plotsense/core/strategies/strategy.py new file mode 100644 index 0000000..52940b6 --- /dev/null +++ b/plotsense/core/strategies/strategy.py @@ -0,0 +1,30 @@ +from abc import ABC, abstractmethod +from typing import List, Tuple, Optional + +class Strategy(ABC): + """ + Base Strategy interface for selecting provider/model pairs. + Each strategy returns a tuple: (provider_name, model_name) + """ + + def __init__(self, provider_models: List[Tuple[str, str]]): + """ + Args: + provider_models: dict mapping provider_name -> list of models + """ + self.provider_models = provider_models + + @abstractmethod + def select_model(self, iteration: int, current_explanation: Optional[str] = None) -> Tuple[str, str]: + """ + Return a (provider, model) tuple for the given iteration. + + Args: + iteration: current iteration index (0-based) + current_explanation: optionally, the current explanation for refinement + + Returns: + (provider_name, model_name) + """ + pass + diff --git a/plotsense/core/utils.py b/plotsense/core/utils.py new file mode 100644 index 0000000..89c9447 --- /dev/null +++ b/plotsense/core/utils.py @@ -0,0 +1,55 @@ +import builtins +import base64 +from matplotlib.figure import Figure +from matplotlib.axes import Axes +from typing import Optional, Union, cast + + +def prompt_for_api_key( + service_name: str, service_link: str, interactive: bool = True, + skip_if_missing: bool = False +) -> Optional[str]: + """Prompt user for API key or raise if unavailable.""" + if not interactive: + if skip_if_missing: + return None + raise ValueError( + f"{service_name.upper()} API key is required. " + f"Set it in the environment or pass it as an argument. " + f"You can get it at {service_link}" + ) + + try: + print(f"βš™οΈ {service_name.upper()} API key not found.") + print(f"πŸ”— Get it at {service_link}") + key = builtins.input(f"Enter {service_name.upper()} API key (or press Enter to skip): ").strip() + if not key and skip_if_missing: + return None + if not key: + raise ValueError(f"{service_name.upper()} API key is required.") + return key + except (EOFError, OSError): + if skip_if_missing: + return None + raise ValueError(f"{service_name.upper()} API key is required (get it at {service_link})") + +def save_plot_to_image( + plot_object: Union[Figure, Axes], + output_path: str = "temp_plot.jpg" +) -> str: + """Save a matplotlib Figure or Axes object to a JPEG image file.""" + if isinstance(plot_object, Axes): + fig = plot_object.figure + else: + fig = plot_object + cast(Figure, fig).savefig( + output_path, format='jpeg', dpi=100, bbox_inches='tight' + ) + return output_path + + +def encode_image(image_path: str) -> str: + """Encode image file to base64 string.""" + with open(image_path, "rb") as image_file: + return base64.b64encode(image_file.read()).decode("utf-8") + diff --git a/plotsense/explanations/explanations.py b/plotsense/explanations/explanations.py index e30533e..46413a1 100644 --- a/plotsense/explanations/explanations.py +++ b/plotsense/explanations/explanations.py @@ -1,186 +1,67 @@ -import base64 import os -import matplotlib.pyplot as plt -from typing import Union, Optional, Dict, List +# import matplotlib.pyplot as plt +from matplotlib.figure import Figure +from matplotlib.axes import Axes +from typing import List, Tuple, Union, Optional, Dict from dotenv import load_dotenv -from groq import Groq -import warnings -import builtins +from plotsense.core.ai_interface import AIModelInterface +from plotsense.core.enums.strategy import StrategyName +from plotsense.core.providers.provider_manager import ProviderManager +from plotsense.core.utils import encode_image, save_plot_to_image load_dotenv() + class PlotExplainer: """ - A class to generate and refine explanations for plots using LLMs.""" - DEFAULT_MODELS = { - 'groq': ['meta-llama/llama-4-scout-17b-16e-instruct', - 'meta-llama/llama-4-maverick-17b-128e-instruct'], - } - + A class to generate and refine explanations for plots using LLMs. + """ + def __init__( - self, - api_keys: Optional[Dict[str, str]] = None, - max_iterations: int = 3, - interactive: bool = True, - timeout: int = 30 + self, + api_keys: Optional[Dict[str, str]], + strategy: StrategyName = StrategyName.ROUND_ROBIN, + selected_models: Optional[List[Tuple[str, str]]] = None, + max_iterations: int = 3, + interactive: bool = True, + timeout: int = 30, ): - # Default to empty dict if None - api_keys = api_keys or {} - - ## Initialize API keys with environment variable or provided keys - self.api_keys = { - 'groq': os.getenv('GROQ_API_KEY') - } - # Update with provided API keys - self.api_keys.update(api_keys) - # Set interactive mode and timeout for API calls - self.interactive = interactive - # Set timeout for API calls - self.timeout = timeout - # Initialize empty dict for clients - self.clients = {} - # Initialize empty list for available models - self.available_models = [] - # Set max iterations for refinement - self.max_iterations = max_iterations - - # Validate API keys and initialize clients - self._validate_keys() - # Initialize clients - self._initialize_clients() - # Detect available models - self._detect_available_models() - - def _validate_keys(self): - """Validate that required API keys are present""" - service_links = { - 'groq': 'πŸ‘‰ https://console.groq.com/keys πŸ‘ˆ' - } - - for service in ['groq']: - if not self.api_keys.get(service): - if self.interactive: - try: - link = service_links.get(service, f"the {service.upper()} website") - message = ( - f"Enter {service.upper()} API key (get it at {link}): " - ) - self.api_keys[service] = builtins.input(message).strip() - if not self.api_keys[service]: - raise ValueError(f"{service.upper()} API key is required") - except (EOFError, OSError): - # Handle cases where input is not available - raise ValueError(f"{service.upper()} API key is required (get it at {service_links.get(service)})") - else: - raise ValueError( - f"{service.upper()} API key is required. " - f"Set it in the environment or pass it as an argument. " - f"You can get it at {service_links.get(service)}" - ) + self.timeout = timeout # timeout for API calls + self.max_iterations = max_iterations # max iterations for refinement + self.strategy_name = strategy # strategy for provider selection - def _initialize_clients(self): - """Initialize API clients based on provided API keys""" - self.clients = {} - if self.api_keys.get('groq'): - try: - self.clients['groq'] = Groq(api_key=self.api_keys['groq']) - except Exception as e: - warnings.warn(f"Could not initialize Groq client: {e}", ImportWarning) - - def _detect_available_models(self): - """Detect available models based on initialized clients""" - self.available_models = [] - for provider, client in self.clients.items(): - if client and provider in self.DEFAULT_MODELS: - self.available_models.extend(self.DEFAULT_MODELS[provider]) - - def save_plot_to_image( - self, - plot_object: Union[plt.Figure, plt.Axes], - output_path: str = "temp_plot.jpg" - ): - """Save plot to an image file""" - if isinstance(plot_object, plt.Axes): - fig = plot_object.figure - else: - fig = plot_object - - fig.savefig(output_path, format='jpeg', dpi=100, bbox_inches='tight') - return output_path - - def encode_image( - self, - image_path: str - ) -> str: - """Encode image file to base64 string""" - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode('utf-8') + selected_providers = {p for p, _ in (selected_models or [])} - def _query_model( - self, - model: str, - prompt: str, - image_path: str, - custom_parameters: Optional[Dict] = None - ) -> str: - - """Generic model querying method with provider-specific logic""" - - base64_image = self.encode_image(image_path) - - # Determine provider based on model name - provider = next( - (p for p, models in self.DEFAULT_MODELS.items() if model in models), - None + self.manager = ProviderManager( + api_keys=api_keys or {}, + interactive=interactive, + restrict_to=list(selected_providers) if selected_providers else None + ) + self.ai_interface = AIModelInterface(self.manager, timeout=self.timeout) + + # if selected_models: + # self.available_models = self.manager.list_all_models + # else: + all_models = self.manager.list_all_models() + self.available_models = [ + (provider, model) + for provider, models in all_models.items() + for model in models + ] + + if not self.available_models: + raise ValueError( + "No available models detected β€” check API keys or selection input." + ) + + self.strategy = self.ai_interface._init_strategy( + self.strategy_name, self.available_models ) - - if not provider: - raise ValueError(f"No provider found for model {model}") - - try: - if provider == 'groq': - client = self.clients['groq'] - - # Merge default and custom parameters - default_params = { - 'max_tokens': 1000, - 'temperature': 0.7 - } - generation_params = {**default_params, **(custom_parameters or {})} - - response = client.chat.completions.create( - model=model, - messages=[ - { - "role": "user", - "content": [ - {"type": "text", "text": prompt}, - { - "type": "image_url", - "image_url": { - "url": f"data:image/jpeg;base64,{base64_image}" - } - } - ] - } - ], - **generation_params - ) - - return response.choices[0].message.content - - except Exception as e: - if "503" in str(e): - print(f"Groq service temporarily unavailable, retrying... Error: {e}") - raise # This will trigger retry - error_message = f"Model querying error for {model}: {str(e)}" - warnings.warn(error_message) - return error_message def refine_plot_explanation( self, - plot_object: Union[plt.Figure, plt.Axes], + plot_object: Union[Figure, Axes], prompt: str = "Explain this data visualization", temp_image_path: str = "temp_plot.jpg", custom_parameters: Optional[Dict] = None @@ -190,40 +71,50 @@ def refine_plot_explanation( raise ValueError("No available models detected") # Save plot to temporary image file - image_path = self.save_plot_to_image(plot_object, temp_image_path) - + image_path = save_plot_to_image(plot_object, temp_image_path) + try: # Iterative refinement process current_explanation = None - + for iteration in range(self.max_iterations): - current_model = self.available_models[iteration % len(self.available_models)] - + provider, current_model = self.strategy.select_model( + iteration, current_explanation + ) + if current_explanation is None: current_explanation = self._generate_initial_explanation( - current_model, image_path, prompt, custom_parameters + provider, current_model, image_path, prompt, custom_parameters ) else: critique = self._generate_critique( - image_path, current_explanation, prompt, current_model, custom_parameters + provider, current_model, image_path, current_explanation, prompt, custom_parameters ) - + current_explanation = self._generate_refinement( - image_path, current_explanation, critique, prompt, current_model, custom_parameters + provider, current_model, image_path, + current_explanation, critique, prompt, + custom_parameters ) + if current_explanation is None: + raise RuntimeError( + "Failed to generate an explanation β€” no models available or initial step failed." + ) + return current_explanation - + finally: # Clean up temporary image file if os.path.exists(image_path): os.remove(image_path) def _generate_initial_explanation( - self, - model: str, + self, + provider: str, + model: str, image_path: str, - original_prompt: str, + original_prompt: str, custom_parameters: Optional[Dict] = None ) -> str: """Generate initial plot explanation with structured format""" @@ -237,7 +128,7 @@ def _generate_initial_explanation( 4. Conclusion - Be specific and data-driven - Highlight key statistical and visual elements - + Specific Prompt: {original_prompt} Formatting Instructions: @@ -246,20 +137,22 @@ def _generate_initial_explanation( - Provide quantitative insights - Explain the significance of visual elements """ - + return self._query_model( - model=model, + provider=provider, + model=model, prompt=base_prompt, - image_path=image_path, + image_path=image_path, custom_parameters=custom_parameters ) def _generate_critique( - self, - image_path: str, - current_explanation: str, - original_prompt: str, + self, + provider: str, model: str, + image_path: str, + current_explanation: str, + original_prompt: str, custom_parameters: Optional[Dict] = None ) -> str: """Generate critique of current explanation""" @@ -293,21 +186,23 @@ def _generate_critique( Provide a constructive critique that will help refine the explanation. """ - + return self._query_model( - model=model, - prompt=critique_prompt, - image_path=image_path, + provider=provider, + model=model, + prompt=critique_prompt, + image_path=image_path, custom_parameters=custom_parameters ) def _generate_refinement( - self, - image_path: str, - current_explanation: str, - critique: str, - original_prompt: str, + self, + provider: str, model: str, + image_path: str, + current_explanation: str, + critique: str, + original_prompt: str, custom_parameters: Optional[Dict] = None ) -> str: """Generate refined explanation based on critique""" @@ -342,50 +237,84 @@ def _generate_refinement( - Use markdown-style headers for clarity - Include bullet points for clarity - Provide quantitative insights - - Ensure the explanation is comprehensive and insightful - + - Ensure the explanation is comprehensive and insightful """ - + return self._query_model( + provider=provider, model=model, - prompt= refinement_prompt, + prompt= refinement_prompt, image_path=image_path, custom_parameters= custom_parameters ) + def _query_model( + self, provider: str, model: str, prompt: str, image_path: str, + custom_parameters: Optional[Dict] = None + ) -> str: + base64_image = encode_image(image_path) + return self.ai_interface.query_model( + provider=provider, + model=model, + prompt=prompt, + base64_image=base64_image, + custom_parameters=custom_parameters + ) + # Package-level convenience function _explainer_instance = None def explainer( - plot_object: Union[plt.Figure, plt.Axes], + plot_object: Union[Figure, Axes], prompt: str = "Explain this data visualization", + *, # force keyword args after this + + custom_parameters: Optional[Dict] = None, + strategy: StrategyName = StrategyName.ROUND_ROBIN, + selected_models: Optional[List[Tuple[str, str]]] = None, + api_keys: Optional[Dict[str, str]] = None, max_iterations: int = 3, - custom_parameters: Optional[Dict] = None, - temp_image_path: str = "temp_plot.jpg" + interactive: bool = True, + timeout: int = 30, + temp_image_path: str = "temp_plot.jpg", ) -> str: """ - Convenience function for iterative plot explanation - + Convenience function to generate and refine plot explanations + Uses a singleton PlotExplainer instance for efficiency. + Args: - data: Original data used to create the plot (DataFrame or numpy array) - plot_object: Matplotlib Figure or Axes - prompt: Explanation prompt - api_keys: API keys for different providers - max_iterations: Maximum refinement iterations - custom_parameters: Additional generation parameters - + - plot_object: Matplotlib Figure or Axes object to explain + - prompt: Initial prompt for explanation generation + - custom_parameters: Optional dict of custom parameters for the model + - strategy: StrategyName enum for model selection strategy + - selected_models: Optional list of (provider, model) tuples to restrict models + - api_keys: Optional dict of API keys for providers + - max_iterations: Max refinement iterations + - interactive: Whether to prompt user for input when needed + - timeout: Timeout in seconds for API calls + - temp_image_path: Path to save temporary plot image + Returns: - Comprehensive explanation with refinement details + A comprehensive, refined explanation generated by the chosen AI models. """ global _explainer_instance + if _explainer_instance is None: - _explainer_instance = PlotExplainer(api_keys=api_keys, - max_iterations=max_iterations) + _explainer_instance = PlotExplainer( + api_keys=api_keys, + strategy=strategy, + selected_models=selected_models, + max_iterations=max_iterations, + interactive=interactive, + timeout=timeout, + ) + return _explainer_instance.refine_plot_explanation( plot_object=plot_object, prompt=prompt, custom_parameters=custom_parameters, temp_image_path=temp_image_path ) + diff --git a/plotsense/plot_generator/__init__.py b/plotsense/plot_generator/__init__.py index 3b22971..ab45af9 100644 --- a/plotsense/plot_generator/__init__.py +++ b/plotsense/plot_generator/__init__.py @@ -1 +1 @@ -from plotsense.plot_generator.generator import plotgen, PlotGenerator \ No newline at end of file +from plotsense.plot_generator.generator import plotgen diff --git a/plotsense/plot_generator/base_generator.py b/plotsense/plot_generator/base_generator.py new file mode 100644 index 0000000..21618b1 --- /dev/null +++ b/plotsense/plot_generator/base_generator.py @@ -0,0 +1,76 @@ +import pandas as pd +from matplotlib.figure import Figure +from typing import Callable, Dict, Optional + +from plotsense.plot_generator.registry import PlotRequirements, PlotTypeRegistry + + +class PlotGenerator: + """ + A class to generate various types of plots based on suggestions. + It uses matplotlib for plotting and can handle both univariate and bivariate cases. + """ + def __init__(self, data, suggestions: Optional[pd.DataFrame] = None): + """ + Initialize with data and plot suggestions. + + Args: + data: DataFrame containing the actual data + suggestions: DataFrame with plot suggestions + """ + if not isinstance(data, pd.DataFrame): + raise TypeError("Data must be a pandas DataFrame") + if data.empty: + raise ValueError("DataFrame is empty") + if not isinstance(suggestions, pd.DataFrame): + raise TypeError("Suggestions must be a pandas DataFrame") + if suggestions.empty: + raise ValueError("Suggestions DataFrame is empty") + if 'plot_type' not in suggestions.columns or 'variables' not in suggestions.columns: + raise ValueError("Suggestions DataFrame must contain 'plot_type' and 'variables' columns") + + self.data = data.copy() + self.suggestions = suggestions + self.registry = PlotTypeRegistry() + self._register_default_plots(self._default_plots) + + @property + def _default_plots(self) -> Dict[str, Callable[..., Figure]]: + """Subclasses override this to define plot type β†’ function mapping.""" + return {} + + def _register_default_plots( + self, plots_to_register: Dict[str, Callable[..., Figure]] + ): + for name, func in plots_to_register.items(): + self.registry.register( + name, + PlotRequirements( + min_variables=1, max_variables=2, numeric_only=True + ), + lambda variables, f=func: f(self.data, variables) + ) + + def generate_plot(self, suggestion_index: int, **kwargs) -> Figure: + """ + Generate a plot based on the suggestion at given index. + + Args: + suggestion_index: Index of the suggestion in dataframe + **kwargs: Additional arguments for the plot + + Returns: + matplotlib Figure object + """ + suggestion = self.suggestions.iloc[suggestion_index] + plot_type = suggestion['plot_type'].lower() + variables = [v.strip() for v in suggestion['variables'].split(',')] + + plot_func = self.registry.get_generator(plot_type) + if not plot_func: + raise ValueError(f"Plot type '{plot_type}' not supported") + + if not self.registry.validate(plot_type, variables, self.data): + raise ValueError(f"Invalid variables for plot '{plot_type}'") + + return plot_func(variables, **kwargs) diff --git a/plotsense/plot_generator/basic_generator.py b/plotsense/plot_generator/basic_generator.py new file mode 100644 index 0000000..7acd6a6 --- /dev/null +++ b/plotsense/plot_generator/basic_generator.py @@ -0,0 +1,29 @@ +from plotsense.plot_generator.base_generator import PlotGenerator +from plotsense.plot_generator.plots.basic.kde import create_kde_plot +from plotsense.plot_generator.plots.basic.barh import create_barh_plot +from plotsense.plot_generator.plots.basic.box import create_box_plot +from plotsense.plot_generator.plots.basic.ecdf import create_ecdf_plot +from plotsense.plot_generator.plots.basic.hist import create_hist_plot +from plotsense.plot_generator.plots.basic.violin import create_violin_plot +from plotsense.plot_generator.plots.basic.bar import create_bar_plot +from plotsense.plot_generator.plots.basic.hexbin import create_hexbin_plot +from plotsense.plot_generator.plots.basic.pie import create_pie_plot +from plotsense.plot_generator.plots.basic.scatter import create_scatter_plot + + +class BasicPlotGenerator(PlotGenerator): + @property + def _default_plots(self): + return { + 'bar': create_bar_plot, + 'barh': create_barh_plot, + 'box': create_box_plot, + 'ecdf': create_ecdf_plot, + 'hexbin': create_hexbin_plot, + 'hist': create_hist_plot, + 'kde': create_kde_plot, + 'pie': create_pie_plot, + 'scatter': create_scatter_plot, + 'violin': create_violin_plot, + } + diff --git a/plotsense/plot_generator/generator.py b/plotsense/plot_generator/generator.py index 54d1c64..721d228 100644 --- a/plotsense/plot_generator/generator.py +++ b/plotsense/plot_generator/generator.py @@ -1,564 +1,132 @@ import pandas as pd -import matplotlib.pyplot as plt -import numpy as np -from typing import List, Dict, Any, Optional, Union - - -class PlotGenerator: - """ - A class to generate various types of plots based on suggestions. - It uses matplotlib for plotting and can handle both univariate and bivariate cases. - """ - def __init__(self, data: pd.DataFrame, suggestions: Optional[pd.DataFrame] = None): - """ - Initialize with data and plot suggestions. - - Args: - data: DataFrame containing the actual data - suggestions: DataFrame with plot suggestions - """ - if not isinstance(data, pd.DataFrame): - raise TypeError("Data must be a pandas DataFrame") - if data.empty: - raise ValueError("DataFrame is empty") - if not isinstance(suggestions, pd.DataFrame): - raise TypeError("Suggestions must be a pandas DataFrame") - if suggestions.empty: - raise ValueError("Suggestions DataFrame is empty") - if 'plot_type' not in suggestions.columns or 'variables' not in suggestions.columns: - raise ValueError("Suggestions DataFrame must contain 'plot_type' and 'variables' columns") - - self.data = data.copy() - self.suggestions = suggestions - self.plot_functions = self._initialize_plot_functions() - - def generate_plot(self, suggestion_index: int, **kwargs) -> plt.Figure: - """ - Generate a plot based on the suggestion at given index. - - Args: - suggestion_index: Index of the suggestion in dataframe - **kwargs: Additional arguments for the plot - - Returns: - matplotlib Figure object - """ - # if suggestion_index < 0 or suggestion_index >= len(self.suggestions): - # raise IndexError("Suggestion index out of range") - if not isinstance(suggestion_index, int): - raise TypeError("Suggestion index must be an integer") - if not isinstance(kwargs, dict): - raise TypeError("Additional arguments must be provided as a dictionary") - if self.suggestions.empty: - raise ValueError("No suggestions available to generate a plot") - if self.data.empty: - raise ValueError("No data available to generate a plot") - if not isinstance(self.suggestions, pd.DataFrame): - raise TypeError("Suggestions must be a pandas DataFrame") - if not isinstance(self.data, pd.DataFrame): - raise TypeError("Data must be a pandas DataFrame") - if self.suggestions.empty: - raise ValueError("Suggestions DataFrame is empty") - if self.data.empty: - raise ValueError("DataFrame is empty") - - suggestion = self.suggestions.iloc[suggestion_index] - plot_type = suggestion['plot_type'].lower() - variables = [v.strip() for v in suggestion['variables'].split(',')] - - if plot_type not in self.plot_functions: - print(f"This version of PlotSense does not support plot type: {plot_type}") - return None - - plot_func = self.plot_functions[plot_type] - return plot_func(variables, **kwargs) - - def _initialize_plot_functions(self) -> Dict[str, callable]: - """Initialize all matplotlib plot functions with their requirements.""" - return { - # Basic plots - 'scatter': self._create_scatter, - 'bar': self._create_bar, - 'barh': self._create_barh, - - # Statistical plots - 'hist': self._create_hist, - 'boxplot': self._create_box, - 'violinplot': self._create_violin, - - # Specialized plots - 'pie': self._create_pie, - 'hexbin': self._create_hexbin - - } - - - # ========== Basic Plot Functions ========== - def _create_scatter(self, variables: List[str], **kwargs) -> plt.Figure: - if len(variables) < 2: - raise ValueError("scatter requires at least 2 variables (x, y)") - fig, ax = plt.subplots() - ax.scatter(self.data[variables[0]], self.data[variables[1]], **kwargs) - self._set_labels(ax, variables) - ax.set_title(f"Scatter: {variables[0]} vs {variables[1]}") - return fig - - def _create_bar(self, variables: List[str], **kwargs) -> plt.Figure: - fig, ax = plt.subplots(figsize=(10, 6)) - - # Extract label-related kwargs if provided - x_label = kwargs.pop('x_label', None) - y_label = kwargs.pop('y_label', None) - title = kwargs.pop('title', None) - - # Define font sizes - tick_fontsize = kwargs.pop('tick_fontsize', 12) - label_fontsize = kwargs.pop('label_fontsize', 14) - title_fontsize = kwargs.pop('title_fontsize', 16) - - if len(variables) == 1: - # Single variable - show value counts - value_counts = self.data[variables[0]].value_counts().sort_values(ascending=False) - ax.bar(value_counts.index.astype(str), value_counts.values, **kwargs) - ax.set_xlabel(variables[0] if x_label is None else x_label, fontsize=label_fontsize) - ax.set_ylabel('Count' if y_label is None else y_label, fontsize=label_fontsize) - ax.set_title(f"Bar plot of {variables[0]}" if title is None else title, fontsize=title_fontsize) - ax.tick_params(axis='x', labelsize=tick_fontsize) - ax.tick_params(axis='y', labelsize=tick_fontsize) - - - if len(value_counts) > 10: - fig.set_size_inches(max(12, len(value_counts)), 8) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - else: - # First variable is numeric, second is categorical - grouped = self.data.groupby(variables[1])[variables[0]].mean().sort_values(ascending=False) - ax.bar(grouped.index.astype(str), grouped.values, **kwargs) - ax.set_xlabel(variables[1] if x_label is None else x_label, fontsize=label_fontsize) - ax.set_ylabel(f"{variables[0]}" if y_label is None else y_label, fontsize=label_fontsize) - ax.set_title(f"{variables[0]} by {variables[1]}" if title is None else title, fontsize=title_fontsize) - ax.tick_params(axis='x', labelsize=tick_fontsize) - ax.tick_params(axis='y', labelsize=tick_fontsize) - - if len(grouped) > 10: - fig.set_size_inches(max(12, len(grouped)), 8) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - return fig - - def _create_barh(self, variables: List[str], **kwargs) -> plt.Figure: - fig, ax = plt.subplots(figsize=(10, 6)) - - - # Extract label-related kwargs if provided - x_label = kwargs.pop('x_label', None) - y_label = kwargs.pop('y_label', None) - title = kwargs.pop('title', None) - - # Define font sizes - tick_fontsize = kwargs.pop('tick_fontsize', 12) - label_fontsize = kwargs.pop('label_fontsize', 14) - title_fontsize = kwargs.pop('title_fontsize', 16) - - if len(variables) == 1: - # Single variable - show value counts - value_counts = self.data[variables[0]].value_counts() - ax.barh(value_counts.index.astype(str), value_counts.values, **kwargs) - ax.set_xlabel(variables[0] if x_label is None else x_label, fontsize=label_fontsize) - ax.set_ylabel('Count' if y_label is None else y_label, fontsize=label_fontsize) - ax.set_title(f"Bar plot of {variables[0]}" if title is None else title, fontsize=title_fontsize) - ax.tick_params(axis='x', labelsize=tick_fontsize) - ax.tick_params(axis='y', labelsize=tick_fontsize) - - - if len(value_counts) > 10: - fig.set_size_inches(max(12, len(value_counts)), 8) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - else: - # First variable is numeric, second is categorical - grouped = self.data.groupby(variables[1])[variables[0]].mean() - ax.barh(grouped.index.astype(str), grouped.values, **kwargs) - ax.set_xlabel(variables[1] if x_label is None else x_label, fontsize=label_fontsize) - ax.set_ylabel(f"{variables[0]}" if y_label is None else y_label, fontsize=label_fontsize) - ax.set_title(f"{variables[0]} by {variables[1]}" if title is None else title, fontsize=title_fontsize) - ax.tick_params(axis='x', labelsize=tick_fontsize) - ax.tick_params(axis='y', labelsize=tick_fontsize) - - if len(grouped) > 10: - fig.set_size_inches(max(12, len(grouped)), 8) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - return fig - - - # ========== Statistical Plot Functions ========== - def _create_hist(self, variables: List[str], **kwargs) -> plt.Figure: - fig, ax = plt.subplots() - ax.hist(self.data[variables[0]], **kwargs) - ax.set_xlabel(variables[0]) - ax.set_ylabel('Frequency') - ax.set_title(f"Histogram of {variables[0]}") - return fig - - def _create_box(self, variables: List[str], **kwargs) -> plt.Figure: - fig, ax = plt.subplots(figsize=(10,6)) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - ax.boxplot(self.data[variables[0]], **kwargs) - ax.set_ylabel(variables[0]) - ax.set_title(f"Box plot of {variables[0]}") - - return fig - - def _create_violin(self, variables: List[str], **kwargs) -> plt.Figure: - fig, ax = plt.subplots(figsize=(10,6)) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - ax.violinplot(self.data[variables[0]], **kwargs) - ax.set_ylabel(variables[0]) - ax.set_title(f"Violin plot of {variables[0]}") - return fig - - - # ========== Specialized Plot Functions ========== - def _create_pie(self, variables: List[str], **kwargs) -> plt.Figure: - value_counts = self.data[variables[0]].value_counts() - fig, ax = plt.subplots() - ax.pie(value_counts, labels=value_counts.index, autopct='%1.1f%%', **kwargs) - ax.set_title(f"Pie chart of {variables[0]}") - return fig - - def _create_hexbin(self, variables: List[str], **kwargs) -> plt.Figure: - fig, ax = plt.subplots() - ax.hexbin(self.data[variables[0]], self.data[variables[1]], **kwargs) - self._set_labels(ax, variables) - ax.set_title(f"Hexbin: {variables[0]} vs {variables[1]}") - return fig - - - - # ========== Helper Methods ========== - def _set_labels(self, ax, variables: List[str]): - """Set labels for x and y axes based on variables.""" - if len(variables) > 0: - ax.set_xlabel(variables[0]) - if len(variables) > 1: - ax.set_ylabel(variables[1]) - -class SmartPlotGenerator(PlotGenerator): - def _create_box(self, variables: List[str], **kwargs) -> plt.Figure: - """Enhanced boxplot that handles both univariate and bivariate cases with NaN handling.""" - fig, ax = plt.subplots(figsize=(10, 6)) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - if len(variables) == 1: - # Univariate case - single numerical variable - data = self.data[variables[0]].dropna() # Remove NaN values - if len(data) == 0: - raise ValueError(f"No valid data remaining after dropping NaN values for {variables[0]}") - ax.boxplot(data, **kwargs) - ax.set_ylabel(variables[0]) - ax.set_title(f"Box plot of {variables[0]}") - elif len(variables) >= 2: - # Bivariate case - numerical vs categorical - numerical_var = variables[0] - categorical_var = variables[1] - - # Clean data - remove rows where either variable is NaN - clean_data = self.data[[numerical_var, categorical_var]].dropna() - if len(clean_data) == 0: - raise ValueError(f"No valid data remaining after cleaning {numerical_var} and {categorical_var}") - - # Group data by categorical variable - grouped_data = [clean_data[clean_data[categorical_var] == cat][numerical_var] - for cat in clean_data[categorical_var].unique()] - - # Filter out empty groups - grouped_data = [group for group in grouped_data if len(group) > 0] - if not grouped_data: - raise ValueError("No valid groups remaining after filtering") - - ax.boxplot(grouped_data, **kwargs) - ax.set_xticklabels(clean_data[categorical_var].unique()) - ax.set_xlabel(categorical_var) - ax.set_ylabel(numerical_var) - ax.set_title(f"Box plot of {numerical_var} by {categorical_var}") - else: - raise ValueError("Box plot requires at least 1 variable") - - return fig - - def _create_violin(self, variables: List[str], **kwargs) -> plt.Figure: - """Enhanced violin plot that handles both univariate and bivariate cases with NaN handling.""" - fig, ax = plt.subplots(figsize=(10,6)) - plt.setp(ax.get_xticklabels(), rotation=90, ha='center') - - if len(variables) == 1: - # Univariate case - single numerical variable - data = self.data[variables[0]].dropna() # Remove NaN values - if len(data) == 0: - raise ValueError(f"No valid data remaining after dropping NaN values for {variables[0]}") - ax.violinplot(data, **kwargs) - ax.set_ylabel(variables[0]) - ax.set_title(f"Violin plot of {variables[0]}") - elif len(variables) >= 2: - # Bivariate case - numerical vs categorical - numerical_var = variables[0] - categorical_var = variables[1] - - # Clean data - remove rows where either variable is NaN - clean_data = self.data[[numerical_var, categorical_var]].dropna() - if len(clean_data) == 0: - raise ValueError(f"No valid data remaining after cleaning {numerical_var} and {categorical_var}") - - # Group data by categorical variable - grouped_data = [clean_data[clean_data[categorical_var] == cat][numerical_var] - for cat in clean_data[categorical_var].unique()] - - # Filter out empty groups - grouped_data = [group for group in grouped_data if len(group) > 0] - if not grouped_data: - raise ValueError("No valid groups remaining after filtering") - - ax.violinplot(grouped_data, **kwargs) - ax.set_xticks(np.arange(1, len(grouped_data)+1)) - ax.set_xticklabels(clean_data[categorical_var].unique()) - ax.set_xlabel(categorical_var) - ax.set_ylabel(numerical_var) - ax.set_title(f"Violin plot of {numerical_var} by {categorical_var}") - else: - raise ValueError("Violin plot requires at least 1 variable") - - return fig - - def _create_hist(self, variables: List[str], **kwargs) -> plt.Figure: - """Enhanced histogram that can handle grouping by a second variable.""" - fig, ax = plt.subplots(figsize=(12, 8)) - - if len(variables) == 1: - # Simple histogram - data = self.data[variables[0]].dropna() - if len(data) == 0: - raise ValueError(f"No valid data remaining for {variables[0]}") - - ax.hist(data, **kwargs) - ax.set_xlabel(variables[0]) - ax.set_ylabel('Frequency') - ax.set_title(f"Histogram of {variables[0]}") - elif len(variables) >= 2: - # Grouped histogram - numerical_var = variables[0] - categorical_var = variables[1] - - # Clean data - clean_data = self.data[[numerical_var, categorical_var]].dropna() - if len(clean_data) == 0: - raise ValueError(f"No valid data remaining after cleaning {numerical_var} and {categorical_var}") - - # Get unique categories - categories = clean_data[categorical_var].unique() - - # Set default colors if not provided - if 'color' not in kwargs and 'colors' not in kwargs: - colors = plt.rcParams['axes.prop_cycle'].by_key()['color'] - else: - colors = [kwargs.pop('color')] * len(categories) if 'color' in kwargs else kwargs.pop('colors') - - # Plot each group - for i, cat in enumerate(categories): - ax.hist(clean_data[clean_data[categorical_var] == cat][numerical_var], - alpha=0.5, - label=str(cat), - color=colors[i % len(colors)], - **kwargs) - - ax.set_xlabel(numerical_var) - ax.set_ylabel('Frequency') - ax.set_title(f"Histogram of {numerical_var} by {categorical_var}") - ax.legend() - else: - raise ValueError("Histogram requires at least 1 variable") - - return fig - - def _create_scatter(self, variables: List[str], - size_scale: float = 100.0, - **kwargs) -> plt.Figure: - """ - Create a scatter plot with optional color and size dimensions. - - Parameters: - ----------- - variables : List[str] - - 2 variables: x, y - - 3 variables: x, y, color - - 4 variables: x, y, color, size - size_scale : float - Scaling factor for bubble sizes (default: 100) - - Returns: - -------- - matplotlib.figure.Figure - """ - if len(variables) < 2: - raise ValueError("Scatter requires at least 2 variables (x, y)") - if len(variables) > 4: - raise ValueError("Scatter supports maximum 4 variables (x, y, color, size)") - - # Check data types - for var in variables[:2]: # x and y must be numeric - if not np.issubdtype(self.data[var].dtype, np.number): - raise ValueError(f"Variable '{var}' must be numeric") - - fig, ax = plt.subplots() - scatter_params = { - 'x': self.data[variables[0]], - 'y': self.data[variables[1]], - } - - # Handle color (3rd variable) - if len(variables) >= 3: - color_data = self.data[variables[2]] - if pd.api.types.is_numeric_dtype(color_data): - # For numeric color data, use continuous colormap - scatter_params['c'] = color_data - kwargs.setdefault('cmap', 'viridis') - else: - # For categorical data, convert to numeric codes - scatter_params['c'] = pd.factorize(color_data)[0] - kwargs.setdefault('cmap', 'tab10') - - # Handle size (4th variable) - if len(variables) == 4: - size_data = self.data[variables[3]] - if not pd.api.types.is_numeric_dtype(size_data): - raise ValueError(f"Size variable '{variables[3]}' must be numeric") - - # Normalize and scale sizes - sizes = np.abs(size_data) # Ensure positive - sizes = (sizes - sizes.min()) / (sizes.max() - sizes.min() + 1e-8) * size_scale - scatter_params['s'] = sizes - - # Apply any additional kwargs - scatter_params.update(kwargs) - - scatter = ax.scatter(**scatter_params) - - # Set labels and title - self._set_labels(ax, variables[:2]) # Assuming this sets x and y labels - title = f"Scatter: {variables[0]} vs {variables[1]}" - if len(variables) >= 3: - title += f" (colored by {variables[2]})" - # Add colorbar for continuous data - if pd.api.types.is_numeric_dtype(self.data[variables[2]]): - fig.colorbar(scatter, ax=ax, label=variables[2]) - if len(variables) == 4: - title += f" (sized by {variables[3]})" - ax.set_title(title) - - return fig - - +from matplotlib.figure import Figure +from typing import Optional, Union, Callable +from plotsense.plot_generator.basic_generator import BasicPlotGenerator +from plotsense.plot_generator.smart_generator import SmartPlotGenerator +from plotsense.plot_generator.registry import PlotRequirements # Global instance of the plot generator _plot_generator_instance = None +_GENERATOR_MAP = { + "basic": BasicPlotGenerator, + "smart": SmartPlotGenerator +} + def plotgen( df: pd.DataFrame, suggestion: Union[int, pd.Series], suggestions_df: Optional[pd.DataFrame] = None, + generator: str = "basic", + plot_function: Optional[Callable] = None, + plot_type: Optional[str] = None, + plot_requirements: Optional[PlotRequirements] = None, **plot_kwargs -) -> plt.Figure: +) -> Figure: """ Generate a plot based on visualization suggestions. - + + Users can also register a custom plot function temporarily by providing: + plot_function: callable(df, variables, **kwargs) -> Figure + plot_type: string name for the custom plot + plot_requirements: optional PlotRequirements object + Args: df: Input DataFrame containing the data to plot suggestion: Either an integer index or a pandas Series containing the suggestion row suggestions_df: DataFrame containing visualization suggestions (required if suggestion is an index) + generator: String identifier for generator to use ("basic" or "smart") + plot_function: Optional custom plot function + plot_type: Name of the custom plot + plot_requirements: Optional PlotRequirements for the custom plot **plot_kwargs: Additional arguments to pass to the plot function - + Returns: matplotlib.Figure: The generated figure - - Example: - # Using index (requires suggestions_df) - fig = plotgen(df, 7, suggestions_df=recommendations) - - # Using direct row access with additional plot arguments - fig = plotgen(df, recommendations.iloc[7], bins=30, color='red') - - # Using specific variable names - fig = plotgen(df, recommendations.iloc[7], x='age', y='fare') """ global _plot_generator_instance - - # Handle case where suggestion is a row from recommendations - if isinstance(suggestion, pd.Series): - # Create a temporary single-row suggestions DataFrame - temp_df = pd.DataFrame([suggestion]) - # Initialize the plot generator with this single suggestion - _plot_generator_instance = SmartPlotGenerator(df, temp_df) - - - # Get the variables from the suggestion - variables = [v.strip() for v in suggestion['variables'].split(',')] - plot_type = suggestion['plot_type'].lower() - - # Handle x, y, z arguments if provided - if 'x' in plot_kwargs: - variables[0] = plot_kwargs.pop('x') - if 'y' in plot_kwargs and len(variables) > 1: - variables[1] = plot_kwargs.pop('y') - if 'z' in plot_kwargs and len(variables) > 2: - variables[2] = plot_kwargs.pop('z') - # Create a new suggestion with updated variables - updated_suggestion = suggestion.copy() - updated_suggestion['variables'] = ','.join(variables) - temp_df = pd.DataFrame([updated_suggestion]) - _plot_generator_instance.suggestions = temp_df - - # Generate the plot - return _plot_generator_instance.generate_plot(0, **plot_kwargs) - - # Handle case where suggestion is an index - elif isinstance(suggestion, int): - if suggestions_df is None: + # Determine generator class from string + generator_class = _GENERATOR_MAP.get(generator.lower(), BasicPlotGenerator) + + # Initialize generator instance if needed + if _plot_generator_instance is None or not isinstance( + _plot_generator_instance, generator_class + ): + # Handle case where suggestion is a row from recommendations + if isinstance(suggestion, pd.Series): + temp_df = pd.DataFrame([suggestion]) + _plot_generator_instance = generator_class(df, temp_df) + # Handle case where suggestion is an index + elif isinstance(suggestion, int): + if suggestions_df is None: + raise ValueError("suggestions_df must be provided when using an index") + _plot_generator_instance = generator_class(df, suggestions_df) + else: + # Update data if it changed + if not _plot_generator_instance.data.equals(df): + _plot_generator_instance.data = df + + # If user provides a custom plot function, register it temporarily + if plot_function is not None: + if not plot_type: + raise ValueError("plot_type name must be provided when registering a custom plot") + if plot_requirements is None: + plot_requirements = PlotRequirements(min_variables=1, max_variables=2, numeric_only=True) + + pg = _plot_generator_instance + if pg is None: + raise RuntimeError("Plot generator instance is not initialized") + + pg.registry.register( + plot_type, + plot_requirements, + lambda variables, + f=plot_function: f(pg.data, variables, **plot_kwargs) + ) + + # Extract suggestion row + if isinstance(suggestion, pd.Series): + suggestion_row = suggestion.copy() + else: + s_df = suggestions_df + if s_df is None: + raise ValueError("suggestions_df must be provided when using an index") + suggestion_row = s_df.iloc[suggestion].copy() + + # Override variables if x/y/z provided + variables = [v.strip() for v in str(suggestion_row['variables']).split(',')] + if 'x' in plot_kwargs: + variables[0] = plot_kwargs.pop('x') + if 'y' in plot_kwargs and len(variables) > 1: + variables[1] = plot_kwargs.pop('y') + if 'z' in plot_kwargs and len(variables) > 2: + variables[2] = plot_kwargs.pop('z') + + suggestion_row['variables'] = ','.join(variables) + + # Update the generator's suggestion DataFrame if using index + if isinstance(suggestion, int): + s_df = suggestions_df + if s_df is None: raise ValueError("suggestions_df must be provided when using an index") - - # Initialize the plot generator if it doesn't exist - if _plot_generator_instance is None: - _plot_generator_instance = SmartPlotGenerator(df, suggestions_df) - else: - # Update the data if the generator exists but the data changed - if not _plot_generator_instance.data.equals(df): - _plot_generator_instance.data = df - - # Get the variables from the suggestion - suggestion_row = suggestions_df.iloc[suggestion] - variables = [v.strip() for v in suggestion_row['variables'].split(',')] - plot_type = suggestion_row['plot_type'].lower() - - # Handle x, y, z arguments if provided - if 'x' in plot_kwargs: - variables[0] = plot_kwargs.pop('x') - if 'y' in plot_kwargs and len(variables) > 1: - variables[1] = plot_kwargs.pop('y') - if 'z' in plot_kwargs and len(variables) > 2: - variables[2] = plot_kwargs.pop('z') - - # Create a new suggestion with updated variables - updated_suggestion = suggestion_row.copy() - updated_suggestion['variables'] = ','.join(variables) - suggestions_df.iloc[suggestion] = updated_suggestion + s_df.iloc[suggestion] = suggestion_row _plot_generator_instance.suggestions = suggestions_df - - # Generate the plot - return _plot_generator_instance.generate_plot(suggestion, **plot_kwargs) - # else: - # raise TypeError("suggestion must be either an integer index or a pandas Series") + else: + _plot_generator_instance.suggestions = pd.DataFrame([suggestion_row]) + + # Determine plot_type to use + active_plot_type = plot_type or str(suggestion_row['plot_type']).lower() + + # Generate the plot + plot_func = _plot_generator_instance.registry.get_generator(active_plot_type) + if not plot_func: + raise ValueError(f"Plot type '{active_plot_type}' not supported") + + if not _plot_generator_instance.registry.validate(active_plot_type, variables, _plot_generator_instance.data): + raise ValueError(f"Invalid variables for plot '{active_plot_type}'") + + return plot_func(variables, **plot_kwargs) +# fig = plotgen(df, 0, suggestions_df, generator="smart") diff --git a/plotsense/plot_generator/helpers.py b/plotsense/plot_generator/helpers.py new file mode 100644 index 0000000..722dd50 --- /dev/null +++ b/plotsense/plot_generator/helpers.py @@ -0,0 +1,10 @@ +from typing import List + + +def set_labels(ax, variables: List[str]): + """Set labels for x and y axes based on variables.""" + if len(variables) > 0: + ax.set_xlabel(variables[0]) + if len(variables) > 1: + ax.set_ylabel(variables[1]) + diff --git a/plotsense/plot_generator/plots/__init__.py b/plotsense/plot_generator/plots/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/plotsense/plot_generator/plots/basic/bar.py b/plotsense/plot_generator/plots/basic/bar.py new file mode 100644 index 0000000..66a9f74 --- /dev/null +++ b/plotsense/plot_generator/plots/basic/bar.py @@ -0,0 +1,46 @@ +from typing import List +import matplotlib.pyplot as plt +import pandas as pd +import numpy as np + +def create_bar_plot(df: pd.DataFrame, variables: List[str], **kwargs): + fig, ax = plt.subplots(figsize=(10, 6)) + + # Extract label-related kwargs if provided + x_label = kwargs.pop('x_label', None) + y_label = kwargs.pop('y_label', None) + title = kwargs.pop('title', None) + + # Define font sizes + tick_fontsize = kwargs.pop('tick_fontsize', 12) + label_fontsize = kwargs.pop('label_fontsize', 14) + title_fontsize = kwargs.pop('title_fontsize', 16) + + if len(variables) == 1: + value_counts = df[variables[0]].value_counts().sort_values(ascending=False) + ax.bar( + value_counts.index.astype(str), + np.asarray(value_counts.values, **kwargs) + ) + ax.set_xlabel(variables[0] if x_label is None else x_label, fontsize=label_fontsize) + ax.set_ylabel('Count' if y_label is None else y_label, fontsize=label_fontsize) + ax.set_title(f"Bar plot of {variables[0]}" if title is None else title, fontsize=title_fontsize) + ax.tick_params(axis='x', labelsize=tick_fontsize) + ax.tick_params(axis='y', labelsize=tick_fontsize) + if len(value_counts) > 10: + fig.set_size_inches(max(12, len(value_counts)), 8) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + else: + grouped = df.groupby(variables[1])[variables[0]].mean() + grouped = pd.Series(grouped).sort_values(ascending=False) + ax.bar(grouped.index.astype(str), np.asarray(grouped.values), **kwargs) + ax.set_xlabel(variables[1] if x_label is None else x_label, fontsize=label_fontsize) + ax.set_ylabel(f"{variables[0]}" if y_label is None else y_label, fontsize=label_fontsize) + ax.set_title(f"{variables[0]} by {variables[1]}" if title is None else title, fontsize=title_fontsize) + ax.tick_params(axis='x', labelsize=tick_fontsize) + ax.tick_params(axis='y', labelsize=tick_fontsize) + if len(grouped) > 10: + fig.set_size_inches(max(12, len(grouped)), 8) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + return fig + diff --git a/plotsense/plot_generator/plots/basic/barh.py b/plotsense/plot_generator/plots/basic/barh.py new file mode 100644 index 0000000..cbf5100 --- /dev/null +++ b/plotsense/plot_generator/plots/basic/barh.py @@ -0,0 +1,48 @@ +from typing import List +from matplotlib.figure import Figure +import matplotlib.pyplot as plt + +def create_barh_plot(df, variables: List[str], **kwargs) -> Figure: + fig, ax = plt.subplots(figsize=(10, 6)) + + + # Extract label-related kwargs if provided + x_label = kwargs.pop('x_label', None) + y_label = kwargs.pop('y_label', None) + title = kwargs.pop('title', None) + + # Define font sizes + tick_fontsize = kwargs.pop('tick_fontsize', 12) + label_fontsize = kwargs.pop('label_fontsize', 14) + title_fontsize = kwargs.pop('title_fontsize', 16) + + if len(variables) == 1: + # Single variable - show value counts + value_counts = df[variables[0]].value_counts() + ax.barh(value_counts.index.astype(str), value_counts.values, **kwargs) + ax.set_xlabel(variables[0] if x_label is None else x_label, fontsize=label_fontsize) + ax.set_ylabel('Count' if y_label is None else y_label, fontsize=label_fontsize) + ax.set_title(f"Bar plot of {variables[0]}" if title is None else title, fontsize=title_fontsize) + ax.tick_params(axis='x', labelsize=tick_fontsize) + ax.tick_params(axis='y', labelsize=tick_fontsize) + + if len(value_counts) > 10: + fig.set_size_inches(max(12, len(value_counts)), 8) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + + else: + # First variable is numeric, second is categorical + grouped = df.groupby(variables[1])[variables[0]].mean() + ax.barh(grouped.index.astype(str), grouped.values, **kwargs) + ax.set_xlabel(variables[1] if x_label is None else x_label, fontsize=label_fontsize) + ax.set_ylabel(f"{variables[0]}" if y_label is None else y_label, fontsize=label_fontsize) + ax.set_title(f"{variables[0]} by {variables[1]}" if title is None else title, fontsize=title_fontsize) + ax.tick_params(axis='x', labelsize=tick_fontsize) + ax.tick_params(axis='y', labelsize=tick_fontsize) + + if len(grouped) > 10: + fig.set_size_inches(max(12, len(grouped)), 8) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + + return fig + diff --git a/plotsense/plot_generator/plots/basic/box.py b/plotsense/plot_generator/plots/basic/box.py new file mode 100644 index 0000000..e9d1733 --- /dev/null +++ b/plotsense/plot_generator/plots/basic/box.py @@ -0,0 +1,15 @@ +from typing import List +from matplotlib.figure import Figure +import matplotlib.pyplot as plt + + +def create_box_plot(df, variables: List[str], **kwargs) -> Figure: + fig, ax = plt.subplots(figsize=(10,6)) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + + ax.boxplot(df[variables[0]], **kwargs) + ax.set_ylabel(variables[0]) + ax.set_title(f"Box plot of {variables[0]}") + + return fig + diff --git a/plotsense/plot_generator/plots/basic/ecdf.py b/plotsense/plot_generator/plots/basic/ecdf.py new file mode 100644 index 0000000..aa20f1f --- /dev/null +++ b/plotsense/plot_generator/plots/basic/ecdf.py @@ -0,0 +1,21 @@ +import matplotlib.pyplot as plt +import numpy as np + +def create_ecdf_plot(df, variables, **kwargs): + """Empirical Cumulative Distribution Function (ECDF) plot.""" + var = variables[0] + data = df[var].dropna() + if data.empty: + raise ValueError(f"No valid data for {var}") + + sorted_data = np.sort(data) + n = len(sorted_data) + y = np.arange(1, n + 1) / n + + fig, ax = plt.subplots(figsize=(8, 5)) + ax.plot(sorted_data, y, marker='.', linestyle='none', **kwargs) + ax.set_title(f"ECDF of {var}") + ax.set_xlabel(var) + ax.set_ylabel("Cumulative Probability") + return fig + diff --git a/plotsense/plot_generator/plots/basic/hexbin.py b/plotsense/plot_generator/plots/basic/hexbin.py new file mode 100644 index 0000000..f574d71 --- /dev/null +++ b/plotsense/plot_generator/plots/basic/hexbin.py @@ -0,0 +1,11 @@ +from typing import List +import matplotlib.pyplot as plt +from matplotlib.figure import Figure + +def create_hexbin_plot(df, variables: List[str], **kwargs) -> Figure: + fig, ax = plt.subplots() + ax.hexbin(df[variables[0]], df[variables[1]], **kwargs) + df._set_labels(ax, variables) + ax.set_title(f"Hexbin: {variables[0]} vs {variables[1]}") + return fig + diff --git a/plotsense/plot_generator/plots/basic/hist.py b/plotsense/plot_generator/plots/basic/hist.py new file mode 100644 index 0000000..c5ea05a --- /dev/null +++ b/plotsense/plot_generator/plots/basic/hist.py @@ -0,0 +1,12 @@ +from typing import List +from matplotlib.figure import Figure +import matplotlib.pyplot as plt + +def create_hist_plot(df, variables: List[str], **kwargs) -> Figure: + fig, ax = plt.subplots() + ax.hist(df[variables[0]], **kwargs) + ax.set_xlabel(variables[0]) + ax.set_ylabel('Frequency') + ax.set_title(f"Histogram of {variables[0]}") + return fig + diff --git a/plotsense/plot_generator/plots/basic/kde.py b/plotsense/plot_generator/plots/basic/kde.py new file mode 100644 index 0000000..88d3c16 --- /dev/null +++ b/plotsense/plot_generator/plots/basic/kde.py @@ -0,0 +1,16 @@ +import matplotlib.pyplot as plt + +def create_kde_plot(df, variables, **kwargs): + """Kernel Density Estimation plot for a numeric variable.""" + var = variables[0] + data = df[var].dropna() + if data.empty: + raise ValueError(f"No valid data for {var}") + + fig, ax = plt.subplots(figsize=(8, 5)) + data.plot(kind='kde', ax=ax, **kwargs) + ax.set_title(f"KDE Plot of {var}") + ax.set_xlabel(var) + ax.set_ylabel("Density") + return fig + diff --git a/plotsense/plot_generator/plots/basic/pie.py b/plotsense/plot_generator/plots/basic/pie.py new file mode 100644 index 0000000..b89e64c --- /dev/null +++ b/plotsense/plot_generator/plots/basic/pie.py @@ -0,0 +1,13 @@ +from typing import List +import matplotlib.pyplot as plt +from matplotlib.figure import Figure + + +def create_pie_plot(df, variables: List[str], **kwargs) -> Figure: + value_counts = df[variables[0]].value_counts() + fig, ax = plt.subplots() + ax.pie(value_counts, labels=value_counts.index, autopct='%1.1f%%', **kwargs) + ax.set_title(f"Pie chart of {variables[0]}") + return fig + + diff --git a/plotsense/plot_generator/plots/basic/scatter.py b/plotsense/plot_generator/plots/basic/scatter.py new file mode 100644 index 0000000..6799dd8 --- /dev/null +++ b/plotsense/plot_generator/plots/basic/scatter.py @@ -0,0 +1,14 @@ +from typing import List +import matplotlib.pyplot as plt + +def create_scatter_plot(df, variables: List[str], **kwargs): + """Scatter plot: requires at least 2 variables (x, y).""" + if len(variables) < 2: + raise ValueError("scatter requires at least 2 variables (x, y)") + fig, ax = plt.subplots() + ax.scatter(df[variables[0]], df[variables[1]], **kwargs) + ax.set_xlabel(variables[0]) + ax.set_ylabel(variables[1]) + ax.set_title(f"Scatter: {variables[0]} vs {variables[1]}") + return fig + diff --git a/plotsense/plot_generator/plots/basic/violin.py b/plotsense/plot_generator/plots/basic/violin.py new file mode 100644 index 0000000..08cfb1a --- /dev/null +++ b/plotsense/plot_generator/plots/basic/violin.py @@ -0,0 +1,14 @@ +from typing import List +from matplotlib.figure import Figure +import matplotlib.pyplot as plt + + +def create_violin_plot(df, variables: List[str], **kwargs) -> Figure: + fig, ax = plt.subplots(figsize=(10,6)) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + + ax.violinplot(df[variables[0]], **kwargs) + ax.set_ylabel(variables[0]) + ax.set_title(f"Violin plot of {variables[0]}") + return fig + diff --git a/plotsense/plot_generator/plots/smart/box.py b/plotsense/plot_generator/plots/smart/box.py new file mode 100644 index 0000000..e45e72c --- /dev/null +++ b/plotsense/plot_generator/plots/smart/box.py @@ -0,0 +1,49 @@ +from typing import List +import matplotlib.pyplot as plt +import pandas as pd +from matplotlib.figure import Figure + +def create_box_plot(df: pd.DataFrame, variables: List[str], **kwargs) -> Figure: + """Enhanced boxplot that handles both univariate and bivariate cases with NaN handling.""" + fig, ax = plt.subplots(figsize=(10, 6)) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + + if len(variables) == 1: + # Univariate case - single numerical variable + data = df[variables[0]].dropna() # Remove NaN values + if data.empty: + raise ValueError(f"No valid data remaining after dropping NaN values for {variables[0]}") + ax.boxplot(data, **kwargs) + ax.set_ylabel(variables[0]) + ax.set_title(f"Box plot of {variables[0]}") + elif len(variables) >= 2: + # Bivariate case - numerical vs categorical + numerical_var = variables[0] + categorical_var = variables[1] + + # Clean data - remove rows where either variable is NaN + clean_data = df[[numerical_var, categorical_var]].dropna() + if clean_data.empty: + raise ValueError(f"No valid data remaining after cleaning {numerical_var} and {categorical_var}") + + # Group data by categorical variable + grouped_data = [ + clean_data[clean_data[categorical_var] == cat][numerical_var] + for cat in clean_data[categorical_var].unique() + ] + + # Filter out empty groups + grouped_data = [group for group in grouped_data if len(group) > 0] + if not grouped_data: + raise ValueError("No valid groups remaining after filtering") + + ax.boxplot(grouped_data, **kwargs) + ax.set_xticklabels(clean_data[categorical_var].unique()) + ax.set_xlabel(categorical_var) + ax.set_ylabel(numerical_var) + ax.set_title(f"Box plot of {numerical_var} by {categorical_var}") + else: + raise ValueError("Box plot requires at least 1 variable") + + return fig + diff --git a/plotsense/plot_generator/plots/smart/ecdf.py b/plotsense/plot_generator/plots/smart/ecdf.py new file mode 100644 index 0000000..09cfe5e --- /dev/null +++ b/plotsense/plot_generator/plots/smart/ecdf.py @@ -0,0 +1,43 @@ +import matplotlib.pyplot as plt +import numpy as np +from typing import List +import pandas as pd +from matplotlib.figure import Figure + +from plotsense.plot_generator.helpers import set_labels + +def create_ecdf_plot(df: pd.DataFrame, variables: List[str], **kwargs) -> Figure: + """ + Enhanced ECDF plot that handles univariate and grouped data with NaN handling. + """ + if len(variables) == 0: + raise ValueError("ECDF plot requires at least 1 variable") + + var = variables[0] + fig, ax = plt.subplots(figsize=(8, 5)) + + if len(variables) == 1: + data = df[var].dropna() + if data.empty: + raise ValueError(f"No valid data for {var}") + sorted_data = np.sort(data) + n = len(sorted_data) + y = np.arange(1, n + 1) / n + ax.plot(sorted_data, y, marker='.', linestyle='none', **kwargs) + else: + # Grouped ECDF + group_var = variables[1] + clean_data = df[[var, group_var]].dropna() + if clean_data.empty: + raise ValueError(f"No valid data after cleaning {var} and {group_var}") + for cat, group in clean_data.groupby(group_var): + sorted_data = np.sort(group[var]) + n = len(sorted_data) + y = np.arange(1, n + 1) / n + ax.plot(sorted_data, y, marker='.', linestyle='none', label=str(cat), **kwargs) + ax.legend(title=group_var) + + ax.set_title(f"ECDF of {var}") + set_labels(ax, variables[:2]) + ax.set_ylabel("Cumulative Probability") + return fig diff --git a/plotsense/plot_generator/plots/smart/histogram.py b/plotsense/plot_generator/plots/smart/histogram.py new file mode 100644 index 0000000..d134320 --- /dev/null +++ b/plotsense/plot_generator/plots/smart/histogram.py @@ -0,0 +1,56 @@ +from typing import List +import matplotlib.pyplot as plt +import pandas as pd +import numpy as np + +def create_histogram_plot(df: pd.DataFrame, variables: List[str], **kwargs) -> plt.Figure: + """Enhanced histogram that can handle grouping by a second variable.""" + fig, ax = plt.subplots(figsize=(12, 8)) + + if len(variables) == 1: + # Simple histogram + data = df[variables[0]].dropna() + if data.empty: + raise ValueError(f"No valid data remaining for {variables[0]}") + ax.hist(data, **kwargs) + ax.set_xlabel(variables[0]) + ax.set_ylabel("Frequency") + ax.set_title(f"Histogram of {variables[0]}") + elif len(variables) >= 2: + # Grouped histogram + num, cat = variables[0], variables[1] + + # Clean data - remove rows where either variable is NaN + clean_data = df[[num, cat]].dropna() + if clean_data.empty: + raise ValueError(f"No valid data remaining after cleaning {num} and {cat}") + + # Get unique categories + categories = clean_data[cat].unique() + + # Set default colors if not provided + if 'color' in kwargs: + colors = [kwargs['color']] * len(categories) + elif 'colors' in kwargs: + colors = kwargs['colors'] + else: + colors = plt.rcParams['axes.prop_cycle'].by_key()['color'] + + # Plot each group + for i, c in enumerate(categories): + ax.hist( + clean_data[clean_data[cat] == c][num], + alpha=0.5, + label=str(c), + color=colors[i % len(colors)], + **kwargs + ) + ax.set_xlabel(num) + ax.set_ylabel("Frequency") + ax.set_title(f"Histogram of {num} by {cat}") + ax.legend() + else: + raise ValueError("Histogram requires at least one variable") + + return fig + diff --git a/plotsense/plot_generator/plots/smart/kde.py b/plotsense/plot_generator/plots/smart/kde.py new file mode 100644 index 0000000..742c361 --- /dev/null +++ b/plotsense/plot_generator/plots/smart/kde.py @@ -0,0 +1,37 @@ +from matplotlib.figure import Figure +import matplotlib.pyplot as plt +import seaborn as sns # optional, for nicer KDE plots +from typing import List +import pandas as pd + +from plotsense.plot_generator.helpers import set_labels + +def create_kde_plot(df: pd.DataFrame, variables: List[str], **kwargs) -> Figure: + """ + Enhanced KDE plot that handles univariate and grouped data with NaN handling. + """ + if len(variables) == 0: + raise ValueError("KDE plot requires at least 1 variable") + + var = variables[0] + data: pd.DataFrame = pd.DataFrame(df[[var]].dropna()) + if data.empty: + raise ValueError(f"No valid data for {var}") + + fig, ax = plt.subplots(figsize=(8, 5)) + + if len(variables) == 1: + # Univariate + sns.kdeplot(data=data, ax=ax, **kwargs) + else: + # Bivariate / group-by + group_var = variables[1] + clean_data: pd.DataFrame = pd.DataFrame(df[[var, group_var]].dropna()) + if clean_data.empty: + raise ValueError(f"No valid data after cleaning {var} and {group_var}") + sns.kdeplot(data=clean_data, x=var, hue=group_var, ax=ax, **kwargs) + + ax.set_title(f"KDE Plot of {var}") + set_labels(ax, variables[:2]) # x + y labels + return fig + diff --git a/plotsense/plot_generator/plots/smart/scatter.py b/plotsense/plot_generator/plots/smart/scatter.py new file mode 100644 index 0000000..6b081cf --- /dev/null +++ b/plotsense/plot_generator/plots/smart/scatter.py @@ -0,0 +1,79 @@ +from typing import List +from matplotlib.figure import Figure +from matplotlib import pyplot as plt +import pandas as pd +import numpy as np + +def create_scatter_plot( + df, variables: List[str], + size_scale: float = 100.0, **kwargs +) -> Figure: + """ + Create a scatter plot with optional color and size dimensions. + + Parameters: + ----------- + variables : List[str] + - 2 variables: x, y + - 3 variables: x, y, color + - 4 variables: x, y, color, size + size_scale : float + Scaling factor for bubble sizes (default: 100) + + Returns: + -------- + matplotlib.figure.Figure + """ + if len(variables) < 2: + raise ValueError("Scatter requires at least 2 variables (x, y)") + if len(variables) > 4: + raise ValueError("Scatter supports up to 4 variables (x, y, color, size)") + + # Check data types + for var in variables[:2]: + if not np.issubdtype(df[var].dtype, np.number): + raise ValueError(f"Variable '{var}' must be numeric") + + fig, ax = plt.subplots() + scatter_params = {"x": df[variables[0]], "y": df[variables[1]]} + + # Handle color (3rd variable) + if len(variables) >= 3: + color_data = df[variables[2]] + if pd.api.types.is_numeric_dtype(color_data): + # For numeric color data, use continuous colormap + scatter_params["c"] = color_data + kwargs.setdefault("cmap", "viridis") + else: + # For categorical data, convert to numeric codes + scatter_params["c"] = pd.factorize(color_data)[0] + kwargs.setdefault("cmap", "tab10") + + # Handle size (4th variable) + if len(variables) == 4: + size_data = df[variables[3]] + if not pd.api.types.is_numeric_dtype(size_data): + raise ValueError(f"Size variable '{variables[3]}' must be numeric") + + # Normalize and scale sizes + sizes = np.abs(size_data) + sizes = (sizes - sizes.min()) / (sizes.max() - sizes.min() + 1e-8) * size_scale + scatter_params["s"] = sizes + + # Apply any additional kwargs + scatter_params.update(kwargs) + scatter = ax.scatter(**scatter_params) + + # Set labels and title + ax.set_xlabel(variables[0]) + ax.set_ylabel(variables[1]) + title = f"Scatter: {variables[0]} vs {variables[1]}" + if len(variables) >= 3: + title += f" (colored by {variables[2]})" + if pd.api.types.is_numeric_dtype(df[variables[2]]): + fig.colorbar(scatter, ax=ax, label=variables[2]) + if len(variables) == 4: + title += f" (sized by {variables[3]})" + ax.set_title(title) + return fig + diff --git a/plotsense/plot_generator/plots/smart/violin.py b/plotsense/plot_generator/plots/smart/violin.py new file mode 100644 index 0000000..1b7b41f --- /dev/null +++ b/plotsense/plot_generator/plots/smart/violin.py @@ -0,0 +1,48 @@ +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +def create_violin_plot(df: pd.DataFrame, variables, **kwargs): + """Enhanced violin plot that handles both univariate and bivariate cases with NaN handling.""" + fig, ax = plt.subplots(figsize=(10, 6)) + plt.setp(ax.get_xticklabels(), rotation=90, ha='center') + + if len(variables) == 1: + # Univariate case - single numerical variable + data = df[variables[0]].dropna() + if data.empty: + raise ValueError(f"No valid data remaining after dropping NaN values for {variables[0]}") + ax.violinplot(data, **kwargs) + ax.set_ylabel(variables[0]) + ax.set_title(f"Violin plot of {variables[0]}") + elif len(variables) >= 2: + # Bivariate case - numerical vs categorical + num, cat = variables[0], variables[1] + + # Clean data - remove rows where either variable is NaN + clean_data = df[[num, cat]].dropna() + if clean_data.empty: + raise ValueError(f"No valid data remaining after cleaning {num} and {cat}") + + # Group data by categorical variable + grouped_data = [ + clean_data[clean_data[cat] == c][num] + for c in clean_data[cat].unique() + ] + + # Filter out empty groups + grouped_data = [g for g in grouped_data if len(g) > 0] + if not grouped_data: + raise ValueError("No valid groups remaining after filtering") + + ax.violinplot(grouped_data, **kwargs) + ax.set_xticks(np.arange(1, len(grouped_data) + 1)) + ax.set_xticklabels(clean_data[cat].unique()) + ax.set_xlabel(cat) + ax.set_ylabel(num) + ax.set_title(f"Violin plot of {num} by {cat}") + else: + raise ValueError("Violin plot requires at least one variable") + + return fig + diff --git a/plotsense/plot_generator/registry.py b/plotsense/plot_generator/registry.py new file mode 100644 index 0000000..f84bf69 --- /dev/null +++ b/plotsense/plot_generator/registry.py @@ -0,0 +1,49 @@ +from typing import Callable, Dict, List, Optional, Any +import pandas as pd +from dataclasses import dataclass + +@dataclass +class PlotRequirements(): + """Define constraints for a plot type.""" + min_variables: int = 1 # minimum required variables + max_variables: int = 2 # maximum supported variables + numeric_only: bool = True # whether data must be numeric + +class PlotTypeRegistry: + """Central registry for all supported plot types.""" + + def __init__(self): + self._registry: Dict[str, Dict[str, Any]] = {} + + def register(self, name: str, requirements: PlotRequirements, generator: Callable): + """Register a plot type and its generation function.""" + self._registry[name.lower()] = { + "requirements": requirements, + "generator": generator + } + + def get_generator(self, name: str) -> Optional[Callable]: + """Retrieve generator function by plot type name.""" + entry = self._registry.get(name.lower()) + return entry["generator"] if entry else None + + def validate(self, name: str, variables: List[str], df: pd.DataFrame) -> bool: + """Check if given data fits the plot type requirements.""" + entry = self._registry.get(name.lower()) + if not entry: + return False + + req = entry["requirements"] + if not (req.min_variables <= len(variables) <= req.max_variables): + return False + + if req.numeric_only: + for var in variables: + if not pd.api.types.is_numeric_dtype(df[var]): + return False + return True + + def list_plot_types(self) -> List[str]: + """List all registered plot types.""" + return list(self._registry.keys()) + diff --git a/plotsense/plot_generator/smart_generator.py b/plotsense/plot_generator/smart_generator.py new file mode 100644 index 0000000..eb83d19 --- /dev/null +++ b/plotsense/plot_generator/smart_generator.py @@ -0,0 +1,25 @@ +from plotsense.plot_generator.base_generator import PlotGenerator +from plotsense.plot_generator.plots.smart.box import create_box_plot +from plotsense.plot_generator.plots.smart.ecdf import create_ecdf_plot +from plotsense.plot_generator.plots.smart.histogram import create_histogram_plot +from plotsense.plot_generator.plots.smart.kde import create_kde_plot +from plotsense.plot_generator.plots.smart.scatter import create_scatter_plot +from plotsense.plot_generator.plots.smart.violin import create_violin_plot + + +class SmartPlotGenerator(PlotGenerator): + """ + An enhanced PlotGenerator with advanced plotting capabilities. + """ + + @property + def _default_plots(self): + return { + 'box': create_box_plot, + 'ecdf': create_ecdf_plot, + 'histogram': create_histogram_plot, + 'kde': create_kde_plot, + 'scatter': create_scatter_plot, + 'violin': create_violin_plot, + } + diff --git a/requirements.txt b/requirements.txt index 343a7b7..54eb62e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,15 +1,149 @@ -# This file is used to install the required packages for the project. -seaborn -matplotlib -pandas -numpy -pytest -python-dotenv -ipykernel -groq -requests -setuptools -jupyter -matplotlib>=3.8.0 -pytest-cov -pytest-mock +annotated-types==0.7.0 +anthropic==0.70.0 +anyio==4.11.0 +argon2-cffi==25.1.0 +argon2-cffi-bindings==25.1.0 +arrow==1.3.0 +asttokens==3.0.0 +async-lru==2.0.5 +attrs==25.4.0 +babel==2.17.0 +beautifulsoup4==4.14.2 +bleach==6.2.0 +cachetools==6.2.1 +certifi==2025.10.5 +cffi==2.0.0 +charset-normalizer==3.4.3 +colorama==0.4.6 +comm==0.2.3 +contourpy==1.3.3 +coverage==7.10.7 +cycler==0.12.1 +debugpy==1.8.17 +decorator==5.2.1 +defusedxml==0.7.1 +distro==1.9.0 +docstring_parser==0.17.0 +executing==2.2.1 +fastjsonschema==2.21.2 +fonttools==4.60.1 +fqdn==1.5.1 +google-ai-generativelanguage==0.6.15 +google-api-core==2.26.0 +google-api-python-client==2.184.0 +google-auth==2.41.1 +google-auth-httplib2==0.2.0 +google-genai==1.45.0 +googleapis-common-protos==1.70.0 +groq==0.32.0 +grpcio==1.75.1 +grpcio-status==1.71.2 +h11==0.16.0 +httpcore==1.0.9 +httplib2==0.31.0 +httpx==0.28.1 +idna==3.10 +iniconfig==2.1.0 +ipykernel==6.30.1 +ipython==9.6.0 +ipython_pygments_lexers==1.1.1 +ipywidgets==8.1.7 +isoduration==20.11.0 +jedi==0.19.2 +Jinja2==3.1.6 +jiter==0.11.0 +json5==0.12.1 +jsonpointer==3.0.0 +jsonschema==4.25.1 +jsonschema-specifications==2025.9.1 +jupyter==1.1.1 +jupyter-console==6.6.3 +jupyter-events==0.12.0 +jupyter-lsp==2.3.0 +jupyter_client==8.6.3 +jupyter_core==5.8.1 +jupyter_server==2.17.0 +jupyter_server_terminals==0.5.3 +jupyterlab==4.4.9 +jupyterlab_pygments==0.3.0 +jupyterlab_server==2.27.3 +jupyterlab_widgets==3.0.15 +kiwisolver==1.4.9 +lark==1.3.0 +MarkupSafe==3.0.3 +matplotlib==3.10.7 +matplotlib-inline==0.1.7 +mistune==3.1.4 +nbclient==0.10.2 +nbconvert==7.16.6 +nbformat==5.10.4 +nest-asyncio==1.6.0 +notebook==7.4.7 +notebook_shim==0.2.4 +numpy==2.3.3 +openai==2.3.0 +packaging==25.0 +pandas==2.3.3 +pandocfilters==1.5.1 +parso==0.8.5 +pillow==11.3.0 +platformdirs==4.5.0 +-e git+https://github.com/DYung26/PlotKit@b5889036ba82bb3cc8c8e50dca1b8238e8bff8f5#egg=plotsense +pluggy==1.6.0 +prometheus_client==0.23.1 +prompt_toolkit==3.0.52 +proto-plus==1.26.1 +protobuf==5.29.5 +psutil==7.1.0 +pure_eval==0.2.3 +pyasn1==0.6.1 +pyasn1_modules==0.4.2 +pycparser==2.23 +pydantic==2.12.0 +pydantic_core==2.41.1 +Pygments==2.19.2 +pyparsing==3.2.5 +pytest==8.4.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.1.1 +python-json-logger==4.0.0 +pytz==2025.2 +pywin32==311 +pywinpty==3.0.2 +PyYAML==6.0.3 +pyzmq==27.1.0 +referencing==0.36.2 +requests==2.32.5 +rfc3339-validator==0.1.4 +rfc3986-validator==0.1.1 +rfc3987-syntax==1.1.0 +rpds-py==0.27.1 +rsa==4.9.1 +seaborn==0.13.2 +Send2Trash==1.8.3 +setuptools==80.9.0 +six==1.17.0 +sniffio==1.3.1 +soupsieve==2.8 +stack-data==0.6.3 +tenacity==9.1.2 +terminado==0.18.1 +tinycss2==1.4.0 +tornado==6.5.2 +tqdm==4.67.1 +traitlets==5.14.3 +types-python-dateutil==2.9.0.20251008 +typing-inspection==0.4.2 +typing_extensions==4.15.0 +tzdata==2025.2 +uri-template==1.3.0 +uritemplate==4.2.0 +urllib3==2.5.0 +wcwidth==0.2.14 +webcolors==24.11.1 +webencodings==0.5.1 +websocket-client==1.9.0 +websockets==15.0.1 +widgetsnbextension==4.0.14 diff --git a/setup.py b/setup.py index 9c244bb..e2a0f52 100644 --- a/setup.py +++ b/setup.py @@ -1,12 +1,17 @@ +# -*- coding: utf-8 -*- +import io from setuptools import setup, find_packages +with io.open("README.md", "r", encoding="utf-8") as f: + long_description = f.read() + setup( name="plotsense", version="0.1.3", author="Christian Chimezie, Toluwaleke Ogidan, Grace Farayola, Amaka Iduwe, Nelson Ogbeide, Onyekachukwu Ojumah, Olamilekan Ajao", author_email="chimeziechristiancc@gmail.com, gbemilekeogidan@gmail.com, gracefarayola@gmail.com, nwaamaka_iduwe@yahoo.com, Ogbeide331@gmail.com, Onyekaojumah22@gmail.com, olamilekan011@gmail.com", description="An intelligent plotting package with suggestions and explanations", - long_description=open("README.md").read(), + long_description=long_description, # open("README.md").read(), long_description_content_type="text/markdown", url="https://github.com/christianchimezie/PlotSenseAI", project_urls={ @@ -34,7 +39,10 @@ "numpy>=1.18", "python-dotenv", "groq", + "anthropic", + "openai", + "google-genai", "requests", ], license="Apache License 2.0", -) \ No newline at end of file +)