edsl 0.1.46__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/{results/DatasetTree.py → dataset/dataset_tree.py}

@@ -136,9 +136,6 @@ class Tree:
             suffix="docx",
             base64_string=base64_string,
         )
-        # doc.save(filename)
-        # from edsl.utilities.utilities import file_notice
-        # file_notice(filename)
 
     def _repr_html_(self):
         """Returns an interactive HTML representation of the tree with collapsible sections."""
@@ -187,8 +184,8 @@ class Tree:
                 content_html = "".join(content)
 
                 if node.children:
-                    # Node with children
-                    html.append(f'<details {"open" if level < 1 else ""}>')
+                    # Node with children - removed the "open" condition to keep all nodes closed by default
+                    html.append("<details>")
                     html.append(f"<summary>{content_html}</summary>")
                     for child in node.children.values():
                         html.append(node_to_html(child, level + 1, print_keys))
@@ -207,40 +204,6 @@ class Tree:
         tree_html = node_to_html(self.root)
         return f"{styles}{tree_html}</div>"
 
-    # def _repr_html_(self):
-    #     """Returns an HTML representation of the tree, following the same logic as print_tree."""
-    #     styles = """
-    #     <style>
-    #         .tree-container {
-    #             font-family: monospace;
-    #             white-space: pre;
-    #             margin: 10px;
-    #         }
-    #     </style>
-    #     """
-
-    #     def node_to_html(node, level=0, print_keys=False):
-    #         if node is None:
-    #             node = self.root
-    #             if node is None:
-    #                 return "Tree has not been constructed yet."
-
-    #         html = []
-    #         if node.value is not None:
-    #             indent = "&nbsp;" * 2 * level  # Using &nbsp; for HTML spaces
-    #             if print_keys and node.key is not None:
-    #                 html.append(f"{indent}{node.key}: {node.value}<br>")
-    #             else:
-    #                 html.append(f"{indent}{node.value}<br>")
-
-    #         for child in node.children.values():
-    #             html.append(node_to_html(child, level + 1, print_keys))
-
-    #         return "".join(html)
-
-    #     tree_html = node_to_html(self.root)
-    #     return f'<div class="tree-container">{tree_html}</div>{styles}'
-
     def _add_to_docx(self, doc, node: TreeNode, level: int):
         if node.value is not None:
             if level == 0:
@@ -256,40 +219,158 @@ class Tree:
         for child in node.children.values():
             self._add_to_docx(doc, child, level + 1)
 
+    def to_dict(self, node: Optional[TreeNode] = None) -> tuple[dict, list[str]]:
+        """Converts the tree structure into a nested dictionary and returns the schema.
+        
+        Args:
+            node: The current node being processed. Defaults to the root node.
+            
+        Returns:
+            A tuple of (nested_dict, schema) where:
+                - nested_dict: The hierarchical data structure
+                - schema: List of keys in order of hierarchy
+            
+        Examples:
+            >>> tree = Tree.example()
+            >>> result, schema = tree.to_dict()
+            >>> print(schema)  # Shows the hierarchy of the data
+            ['continent', 'country', 'city', 'population']
+            >>> # Access a leaf node value (note: numbers are converted to strings)
+            >>> print(result['North America']['US']['New York']['8419000'] is None)
+            True
+            >>> # Verify the structure is correct
+            >>> print(sorted(result['North America'].keys()))
+            ['Canada', 'US']
+        """
+        if node is None:
+            node = self.root
+            if node is None:
+                return {}, []
+
+        result = {}
+        for value, child in node.children.items():
+            if child.children:
+                nested_dict, _ = self.to_dict(child)  # Always unpack tuple, ignore schema except at root
+                result[value] = nested_dict
+            else:
+                result[value] = None
+
+        # Only return the schema with the root call
+        if node == self.root:
+            return result, self.node_order if self.node_order else []
+        return result, []  # Return empty schema for non-root nodes
+
+    @classmethod
+    def example(cls) -> "Tree":
+        """Creates an example Tree instance with geographic data.
+        
+        Returns:
+            Tree: A sample tree with continent/country/city/population data
+        
+        Examples:
+            >>> tree = Tree.example()
+            >>> result, schema = tree.to_dict()
+            >>> print(schema)
+            ['continent', 'country', 'city', 'population']
+            >>> # Verify the structure is correct
+            >>> print(sorted(result['North America'].keys()))
+            ['Canada', 'US']
+        """
+        from .dataset import Dataset
+        
+        data = Dataset([
+            {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
+            {"country": ["US", "China", "France", "Canada", "Japan"]},
+            {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
+            {"population": ["8419000", "21540000", "2161000", "2930000", "13960000"]},  # Convert to strings
+        ])
+        
+        node_order = ["continent", "country", "city", "population"]
+        tree = cls(data, node_order=node_order)  # Explicitly pass node_order
+        return tree
+
+    @staticmethod
+    def _adjust_markdown_levels(text: str, base_level: int) -> str:
+        """Adjusts markdown heading levels by adding or removing '#' characters.
+        
+        Args:
+            text: The markdown text to adjust
+            base_level: The level to adjust headings to (e.g., 2 means h2)
+            
+        Returns:
+            Adjusted markdown text with updated heading levels
+            
+        Examples:
+            >>> text = "# Title\\n## Subtitle\\nContent"
+            >>> print(Tree._adjust_markdown_levels(text, 2))
+            ## Title
+            ### Subtitle
+            Content
+        """
+        lines = []
+        for line in text.split('\n'):
+            if line.strip().startswith('#'):
+                # Count leading '#' characters
+                heading_level = len(line) - len(line.lstrip('#'))
+                # Adjust the heading level by adding base_level - 1
+                new_level = heading_level + (base_level - 1)
+                # Replace the original heading markers with the new level
+                lines.append('#' * new_level + line[heading_level:])
+            else:
+                lines.append(line)
+        return '\n'.join(lines)
+
+    def report(self, node: Optional[TreeNode] = None, level: int = 1, render: bool = True) -> str:
+        """Generates a markdown document representing the tree structure.
+        
+        Args:
+            node: The current node being processed. Defaults to the root node.
+            level: Current heading level (h1-h6). Defaults to 1.
+            render: Whether to render the markdown in notebooks. Defaults to True.
+            
+        Returns:
+            A string containing the markdown document, or renders markdown in notebooks.            
+        """
+        from edsl.utilities.utilities import is_notebook
+        from IPython.display import Markdown, display
+        
+        if node is None:
+            node = self.root
+            if node is None:
+                return "Tree has not been constructed yet."
 
-# Example usage (commented out)
-"""
-from edsl.results.Dataset import Dataset
-
-data = Dataset(
-    [
-        {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
-        {"country": ["US", "China", "France", "Canada", "Japan"]},
-        {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
-        {"population": [8419000, 21540000, 2161000, 2930000, 13960000]},
-    ]
-)
-
-tree = Tree(data)
-
-try:
-    tree.construct_tree(["continent", "country", "city", "population"])
-    print("Tree without key names:")
-    tree.print_tree()
-    print("\nTree with key names:")
-    tree.print_tree(print_keys=True)
-except ValueError as e:
-    print(f"Error: {e}")
-
-# Demonstrating validation
-try:
-    tree.construct_tree(["continent", "country", "invalid_key"])
-except ValueError as e:
-    print(f"\nValidation Error: {e}")
-
-tree = Tree(data)
-tree.construct_tree(["continent", "country", "city", "population"])
-tree.print_tree(print_keys=True)
-tree.to_docx("tree_structure.docx")
-print("DocX file 'tree_structure.docx' has been created.")
-"""
+        lines = []
+        
+        # Process current node
+        if node != self.root:  # Skip the root node as it has no value
+            if isinstance(node.value, str) and node.value.startswith('#'):
+                # If the value is markdown, adjust its heading levels
+                adjusted_markdown = self._adjust_markdown_levels(node.value, level)
+                lines.append(adjusted_markdown)
+            elif node.children:  # Non-leaf nodes get headings
+                # Ensure we don't exceed h6
+                heading_level = min(level, 6)
+                lines.append(f"{'#' * heading_level} {node.key.title()}: {node.value}")
+            else:  # Leaf nodes get regular text
+                lines.append(f"{node.key.title()}: {node.value}")
+        
+        # Process children in sorted order for consistent output
+        for value in sorted(node.children.keys()):
+            child = node.children[value]
+            lines.append(self.report(child, level + 1, render=False))  # Don't render recursive calls
+        
+        markdown_text = "\n".join(lines)
+        
+        # Only attempt to render at the top level call
+        if node == self.root and render and is_notebook():
+            display(Markdown(markdown_text))
+            return ""  # Return empty string since we've displayed the content
+        
+        return markdown_text
+
+
+if __name__ == "__main__":
+    #tree = Tree.example()
+    #print(tree.to_dict())
+    import doctest 
+    doctest.testmod()

edsl/{results/TableDisplay.py → dataset/display/table_display.py}

@@ -10,15 +10,15 @@ from typing import (
 )
 
 if TYPE_CHECKING:
-    from edsl.results.Dataset import Dataset
+    from ..dataset import Dataset
 
-from edsl.results.table_data_class import TableData
+from .table_data_class import TableData
 
-from edsl.results.table_renderers import DataTablesRenderer, PandasStyleRenderer
+from .table_renderers import DataTablesRenderer, PandasStyleRenderer, RichRenderer
 
 Row = Sequence[Union[str, int, float, bool, None]]
 TableFormat = Literal[
-    "grid", "simple", "pipe", "orgtbl", "rst", "mediawiki", "html", "latex"
+    "grid", "simple", "pipe", "orgtbl", "rst", "mediawiki", "html", "latex", "rich"
 ]
 
 
@@ -66,9 +66,20 @@ class TableDisplay:
         return self.renderer_class(table_data).render_html()
 
     def __repr__(self):
-        from tabulate import tabulate
-
-        return tabulate(self.data, headers=self.headers, tablefmt=self.tablefmt)
+        # If rich format is requested, use RichRenderer
+        if self.tablefmt == "rich":
+            table_data = TableData(
+                headers=self.headers,
+                data=self.data,
+                parameters=self.printing_parameters,
+                raw_data_set=self.raw_data_set,
+            )
+            RichRenderer(table_data).render_terminal()
+            return ""  # Return empty string since the table is already printed
+        else:
+            # Fall back to tabulate for other formats
+            from tabulate import tabulate
+            return tabulate(self.data, headers=self.headers, tablefmt=self.tablefmt)
 
     @classmethod
     def from_dictionary(

edsl/{results → dataset/display}/table_renderers.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
-from edsl.results.table_data_class import TableData
-
+import os
+from pathlib import Path
+from .table_data_class import TableData
 
 class DataTablesRendererABC(ABC):
     def __init__(self, table_data: TableData):
@@ -116,3 +117,58 @@ class PandasStyleRenderer(DataTablesRendererABC):
     @classmethod
     def get_css(cls) -> str:
         return ""  # Pandas styling handles its own CSS
+        
+        
+class RichRenderer(DataTablesRendererABC):
+    """Rich-based terminal renderer implementation"""
+    
+    def render_html(self) -> str:
+        """
+        This method is required by the ABC but is not the primary function
+        for this renderer. The render_terminal method below is what's used.
+        """
+        # Provide a basic HTML fallback for HTML contexts
+        html = "<table border='1'><thead><tr>"
+        html += "".join(f"<th>{header}</th>" for header in self.table_data.headers)
+        html += "</tr></thead><tbody>"
+        
+        for row in self.table_data.data:
+            html += "<tr>"
+            html += "".join(f"<td>{cell}</td>" for cell in row)
+            html += "</tr>"
+        html += "</tbody></table>"
+        
+        return html
+    
+    def render_terminal(self) -> None:
+        """
+        Render the table to the terminal using Rich Console and Table
+        
+        This is the primary rendering method for this renderer.
+        """
+        try:
+            from rich.console import Console
+            from rich.table import Table
+            
+            # Create a table with the headers
+            table = Table(show_header=True, header_style="bold")
+            
+            # Add columns with headers
+            for header in self.table_data.headers:
+                table.add_column(str(header))
+            
+            # Add rows
+            for row in self.table_data.data:
+                # Convert all values to strings and handle None
+                str_row = [str(cell) if cell is not None else "" for cell in row]
+                table.add_row(*str_row)
+            
+            # Create a console and print the table
+            console = Console()
+            console.print(table)
+            
+        except ImportError:
+            # Fallback if rich is not installed
+            print("Rich package is not installed. Install with 'pip install rich'")
+            from tabulate import tabulate
+            print(tabulate(self.table_data.data, headers=self.table_data.headers, tablefmt="grid"))

edsl/{results → dataset}/file_exports.py

@@ -3,9 +3,6 @@ import io
 import csv
 import base64
 from typing import Optional, Union, Tuple, List, Any, Dict
-from openpyxl import Workbook
-
-from edsl.scenarios.FileStore import FileStore
 
 
 class FileExport(ABC):
@@ -47,7 +44,7 @@ class FileExport(ABC):
         else:
             base64_string = base64.b64encode(data).decode()
 
-        from edsl.scenarios.FileStore import FileStore
+        from edsl.scenarios import FileStore
 
         path = self.filename or self._get_default_filename()
 
@@ -93,7 +90,7 @@ class TabularExport(FileExport, ABC):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.header, self.rows = self.data._get_tabular_data(
+        self.header, self.rows = self.data.get_tabular_data(
             remove_prefix=self.remove_prefix, pretty_labels=self.pretty_labels
         )
 
@@ -121,6 +118,8 @@ class ExcelExport(TabularExport):
         self.sheet_name = sheet_name or "Results"
 
     def format_data(self) -> bytes:
+        from openpyxl import Workbook
+
         wb = Workbook()
         ws = wb.active
         ws.title = self.sheet_name

edsl/{results → dataset}/smart_objects.py

@@ -35,13 +35,13 @@ class SmartStr(str):
 
     def pdf(self, filename: Optional[str] = None):  # Markdown will have this as well
         # renders the markdown as a pdf that can be downloaded
-        from edsl.results.MarkdownToPDF import MarkdownToPDF
+        from ..utilities.markdown_to_pdf import MarkdownToPDF
 
         return MarkdownToPDF(self, filename).preview()
 
     def docx(self, filename: Optional[str] = None):
         # renders the markdown as a docx that can be downloaded
-        from edsl.results.MarkdownToDocx import MarkdownToDocx
+        from ..utilities.markdown_to_docx import MarkdownToDocx
 
         return MarkdownToDocx(self, filename).preview()
 

edsl/enums.py

@@ -1,207 +1,7 @@
-"""Enums for the different types of questions, language models, and inference services."""
+"""Enums module wrapper for backward compatibility.
 
-from enum import Enum
-from typing import Literal
+This module re-exports everything from the edsl.base.enums module.
+"""
 
-
-class EnumWithChecks(Enum):
-    """Base class for all enums with checks."""
-
-    @classmethod
-    def is_value_valid(cls, value):
-        """Check if the value is valid."""
-        return any(value == item.value for item in cls)
-
-
-class QuestionType(EnumWithChecks):
-    """Enum for the question types."""
-
-    MULTIPLE_CHOICE = "multiple_choice"
-    YES_NO = "yes_no"
-    FREE_TEXT = "free_text"
-    RANK = "rank"
-    BUDGET = "budget"
-    CHECKBOX = "checkbox"
-    EXTRACT = "extract"
-    FUNCTIONAL = "functional"
-    LIST = "list"
-    NUMERICAL = "numerical"
-    TOP_K = "top_k"
-    LIKERT_FIVE = "likert_five"
-    LINEAR_SCALE = "linear_scale"
-
-
-# https://huggingface.co/meta-llama/Llama-2-70b-chat-hf
-
-
-# class LanguageModelType(EnumWithChecks):
-#     """Enum for the language model types."""
-
-#     GPT_4 = "gpt-4-1106-preview"
-#     GPT_3_5_Turbo = "gpt-3.5-turbo"
-#     LLAMA_2_70B_CHAT_HF = "llama-2-70b-chat-hf"
-#     LLAMA_2_13B_CHAT_HF = "llama-2-13b-chat-hf"
-#     GEMINI_PRO = "gemini_pro"
-#     MIXTRAL_8x7B_INSTRUCT = "mixtral-8x7B-instruct-v0.1"
-#     TEST = "test"
-#     ANTHROPIC_3_OPUS = "claude-3-opus-20240229"
-#     ANTHROPIC_3_SONNET = "claude-3-sonnet-20240229"
-#     ANTHROPIC_3_HAIKU = "claude-3-haiku-20240307"
-#     DBRX_INSTRUCT = "dbrx-instruct"
-
-
-class InferenceServiceType(EnumWithChecks):
-    """Enum for the inference service types."""
-
-    BEDROCK = "bedrock"
-    DEEP_INFRA = "deep_infra"
-    REPLICATE = "replicate"
-    OPENAI = "openai"
-    GOOGLE = "google"
-    TEST = "test"
-    ANTHROPIC = "anthropic"
-    GROQ = "groq"
-    AZURE = "azure"
-    OLLAMA = "ollama"
-    MISTRAL = "mistral"
-    TOGETHER = "together"
-    PERPLEXITY = "perplexity"
-    DEEPSEEK = "deepseek"
-    XAI = "xai"
-
-
-# unavoidable violation of the DRY principle but it is necessary
-# checked w/ a unit test to make sure consistent with services in enums.py
-InferenceServiceLiteral = Literal[
-    "bedrock",
-    "deep_infra",
-    "replicate",
-    "openai",
-    "google",
-    "test",
-    "anthropic",
-    "groq",
-    "azure",
-    "ollama",
-    "mistral",
-    "together",
-    "perplexity",
-    "deepseek",
-    "xai",
-]
-
-available_models_urls = {
-    "anthropic": "https://docs.anthropic.com/en/docs/about-claude/models",
-    "openai": "https://platform.openai.com/docs/models/gp",
-    "groq": "https://console.groq.com/docs/models",
-    "google": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models",
-}
-
-
-service_to_api_keyname = {
-    InferenceServiceType.DEEP_INFRA.value: "DEEP_INFRA_API_KEY",
-    InferenceServiceType.REPLICATE.value: "TBD",
-    InferenceServiceType.OPENAI.value: "OPENAI_API_KEY",
-    InferenceServiceType.GOOGLE.value: "GOOGLE_API_KEY",
-    InferenceServiceType.TEST.value: "TBD",
-    InferenceServiceType.ANTHROPIC.value: "ANTHROPIC_API_KEY",
-    InferenceServiceType.GROQ.value: "GROQ_API_KEY",
-    InferenceServiceType.BEDROCK.value: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
-    InferenceServiceType.MISTRAL.value: "MISTRAL_API_KEY",
-    InferenceServiceType.TOGETHER.value: "TOGETHER_API_KEY",
-    InferenceServiceType.PERPLEXITY.value: "PERPLEXITY_API_KEY",
-    InferenceServiceType.DEEPSEEK.value: "DEEPSEEK_API_KEY",
-    InferenceServiceType.XAI.value: "XAI_API_KEY",
-}
-
-
-class TokenPricing:
-    def __init__(
-        self,
-        *,
-        model_name,
-        prompt_token_price_per_k: float,
-        completion_token_price_per_k: float,
-    ):
-        self.model_name = model_name
-        self.prompt_token_price = prompt_token_price_per_k / 1_000.0
-        self.completion_token_price = completion_token_price_per_k / 1_000.0
-
-    def __eq__(self, other):
-        if not isinstance(other, TokenPricing):
-            return False
-        return (
-            self.model_name == other.model_name
-            and self.prompt_token_price == other.prompt_token_price
-            and self.completion_token_price == other.completion_token_price
-        )
-
-
-pricing = {
-    "dbrx-instruct": TokenPricing(
-        model_name="dbrx-instruct",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "claude-3-opus-20240229": TokenPricing(
-        model_name="claude-3-opus-20240229",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "claude-3-haiku-20240307": TokenPricing(
-        model_name="claude-3-haiku-20240307",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "claude-3-sonnet-20240229": TokenPricing(
-        model_name="claude-3-sonnet-20240229",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "gpt-3.5-turbo": TokenPricing(
-        model_name="gpt-3.5-turbo",
-        prompt_token_price_per_k=0.0005,
-        completion_token_price_per_k=0.0015,
-    ),
-    "gpt-4-1106-preview": TokenPricing(
-        model_name="gpt-4",
-        prompt_token_price_per_k=0.01,
-        completion_token_price_per_k=0.03,
-    ),
-    "test": TokenPricing(
-        model_name="test",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "gemini_pro": TokenPricing(
-        model_name="gemini_pro",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "llama-2-13b-chat-hf": TokenPricing(
-        model_name="llama-2-13b-chat-hf",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "llama-2-70b-chat-hf": TokenPricing(
-        model_name="llama-2-70b-chat-hf",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "mixtral-8x7B-instruct-v0.1": TokenPricing(
-        model_name="mixtral-8x7B-instruct-v0.1",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-}
-
-
-def get_token_pricing(model_name):
-    if model_name in pricing:
-        return pricing[model_name]
-    else:
-        return TokenPricing(
-            model_name=model_name,
-            prompt_token_price_per_k=0.0,
-            completion_token_price_per_k=0.0,
-        )
+# Re-export everything from edsl.base.enums module
+from edsl.base.enums import *

edsl/inference_services/__init__.py

@@ -0,0 +1,5 @@
+from .inference_services_collection import InferenceServicesCollection
+from .registry import default
+from .inference_service_abc import InferenceServiceABC
+from .data_structures import AvailableModels
+from .exceptions import InferenceServiceError

edsl/inference_services/{AvailableModelCacheHandler.py → available_model_cache_handler.py}

@@ -4,11 +4,10 @@ import sqlite3
 from datetime import datetime
 import tempfile
 from platformdirs import user_cache_dir
-from dataclasses import dataclass
 import os
 
-from edsl.inference_services.data_structures import LanguageModelInfo, AvailableModels
-from edsl.enums import InferenceServiceLiteral
+from .data_structures import LanguageModelInfo, AvailableModels
+from ..enums import InferenceServiceLiteral
 
 
 class AvailableModelCacheHandler: