pop-python 1.0.3__py3-none-any.whl → 1.1.0__py3-none-any.whl

POP/Embedder.py

@@ -1,44 +1,74 @@
-# Embedder.py
+"""
+Embedding utilities for POP.
+
+This module implements a unified embedding interface capable of
+fetching embeddings via third‑party APIs (JinaAI, OpenAI) or via
+a local PyTorch model.  It is largely derived from the original
+POP project’s ``Embedder.py`` and can be used independently of
+``PromptFunction``.
+
+Example usage:
+
+>>> from pop.embedder import Embedder
+>>> embedder = Embedder(use_api='openai')
+>>> vectors = embedder.get_embedding(["Hello, world!"])
+
+The return value is a numpy array of shape (n_texts, embedding_dim).
+"""
+
 import numpy as np
 import openai
-import requests as HTTPRequests ## some packages already have "requests"
+import requests as HTTPRequests
 from os import getenv
 from backoff import on_exception, expo
-
+from typing import List
 
 from transformers import AutoTokenizer, AutoModel
 
+# Maximum number of tokens permitted by the Jina segmenter
 MAX_TOKENS = 8194
 
 class Embedder:
-    def __init__(self, model_name=None, use_api=None, to_cuda=False, attn_implementation=None):
-        """
-        Initializes the Embedder class, which supports multiple embedding methods, including Jina API, 
-        OpenAI API, and local model embeddings.
-        
-        Args:
-            model_name (str): Name of the model to use for embedding.
-            use_api (str): Flag to determine whether to use an API for embedding ('jina', 'openai') or a local model (None).
-            to_cuda (bool): If True, use GPU; otherwise use CPU. (Some model must run on GPU)
-            attn_implementation (str): Attention implementation method for the transformer model.
-        """
+    """
+    A class supporting multiple embedding methods, including Jina API,
+    OpenAI API, and local model embeddings via PyTorch.
+
+    Parameters
+    ----------
+    model_name:
+        Name of the model to use for embedding.  If ``None`` the default
+        model for the selected API will be chosen.
+    use_api:
+        Which API to use for embedding.  Supported values are
+        ``'jina'``, ``'openai'`` and ``None`` (for local embedding).
+    to_cuda:
+        If ``True``, use GPU; otherwise use CPU for local embeddings.
+    attn_implementation:
+        Optional attention implementation to pass to the transformer
+        when loading the local model.
+    """
+
+    def __init__(self, model_name: str = None, use_api: str = None,
+                 to_cuda: bool = False, attn_implementation: str = None):
         self.use_api = use_api
         self.model_name = model_name
         self.to_cuda = to_cuda
 
-        # API-based embedding initialization
-        if self.use_api or self.use_api == "":
-            supported_apis = ["", 'jina', 'openai',]
+        # API‑based embedding initialisation
+        if self.use_api is not None:
+            supported_apis = ['', 'jina', 'openai']
             if self.use_api not in supported_apis:
                 raise ValueError(f"API type '{self.use_api}' not supported. Supported APIs: {supported_apis}")
-            
-            elif self.use_api == "": # default
-                self.use_api == 'openai'
 
-            elif self.use_api == 'jina':
-                pass # maybe add something later
+            if self.use_api == '':
+                # empty string falls back to OpenAI
+                self.use_api = 'openai'
 
+            if self.use_api == 'jina':
+                # The Jina client requires an API key; nothing to initialise
+                self.client = None
             elif self.use_api == 'openai':
+                # Initialise OpenAI client
                 self.client = openai.Client(api_key=getenv("OPENAI_API_KEY"))
         else:
             # Load PyTorch model for local embedding generation
@@ -47,92 +77,86 @@ class Embedder:
             self.attn_implementation = attn_implementation
             self._initialize_local_model()
 
-    def _initialize_local_model(self):
-        import torch  # Importing PyTorch only when needed
+    def _initialize_local_model(self) -> None:
+        """Initialise the PyTorch model and tokenizer for local embedding generation."""
+        import torch
         import torch.nn.functional as F
 
-
-        """Initializes the PyTorch model and tokenizer for local embedding generation."""
         if self.attn_implementation:
-            self.model = AutoModel.from_pretrained(self.model_name, 
-                                                   trust_remote_code=True, 
-                                                   attn_implementation=self.attn_implementation, 
-                                                   torch_dtype=torch.float16).to('cuda' if self.to_cuda else 'cpu')
+            self.model = AutoModel.from_pretrained(
+                self.model_name,
+                trust_remote_code=True,
+                attn_implementation=self.attn_implementation,
+                torch_dtype=torch.float16,
+            ).to('cuda' if self.to_cuda else 'cpu')
         else:
-            self.model = AutoModel.from_pretrained(self.model_name, 
-                                                   trust_remote_code=True, 
-                                                   torch_dtype=torch.float16).to('cuda' if self.to_cuda else 'cpu')
+            self.model = AutoModel.from_pretrained(
+                self.model_name,
+                trust_remote_code=True,
+                torch_dtype=torch.float16,
+            ).to('cuda' if self.to_cuda else 'cpu')
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         self.model.eval()
 
-    def get_embedding(self, texts: list) -> np.ndarray:
-        """
-        Generates embeddings for a list of texts.
-        
-        Args:
-            texts (list of str): A list of texts to be embedded.
-        
-        Returns:
-            np.ndarray: The embeddings as a numpy array of shape (len(texts), embedding_dim).
+    def get_embedding(self, texts: List[str]) -> np.ndarray:
         """
+        Generate embeddings for a list of texts.
 
+        Parameters
+        ----------
+        texts:
+            A list of strings to embed.
+
+        Returns
+        -------
+        numpy.ndarray
+            Embeddings as a 2‑D array of shape (len(texts), embedding_dim).
+        """
         if not isinstance(texts, list):
             raise ValueError("Input must be a list of strings.")
-        
 
         if self.use_api:
             if self.use_api == 'jina':
+                # set default model if not provided
                 if not self.model_name:
                     self.model_name = "jina-embeddings-v3"
-                    print(f"use default model: {self.model_name}")
                 return self._get_jina_embedding(texts)
             elif self.use_api == 'openai':
-                # set the default to be GPT embedding
                 if not self.model_name:
                     self.model_name = "text-embedding-3-small"
-                    print(f"use default model: {self.model_name}")
                 return self._get_openai_embedding(texts)
             else:
                 raise ValueError(f"API type '{self.use_api}' is not supported.")
         else:
             return self._get_torch_embedding(texts)
-    
-    ## Below are model-specific functions
 
     @on_exception(expo, HTTPRequests.exceptions.RequestException, max_time=30)
-    def _get_jina_embedding(self, texts: list) -> np.ndarray:
-        """Fetches embeddings from the Jina API. Requires Jina API key in .env file."""
+    def _get_jina_embedding(self, texts: List[str]) -> np.ndarray:
+        """Fetch embeddings from the Jina API.  Requires Jina API key in .env."""
         url = 'https://api.jina.ai/v1/embeddings'
-
         headers = {
             'Content-Type': 'application/json',
-            'Authorization': f'Bearer {getenv("JINAAI_API_KEY")}'
+            'Authorization': f"Bearer {getenv('JINAAI_API_KEY')}"
         }
-
-        input_texts = [text for text in texts]
         data = {
-            "model": "jina-embeddings-v3",
+            "model": self.model_name or "jina-embeddings-v3",
             "task": "text-matching",
             "dimensions": 1024,
             "late_chunking": False,
             "embedding_type": "float",
-            "input": input_texts
+            "input": [text for text in texts],
         }
-        response = HTTPRequests.post(url, headers=headers, json=data)    
-
-        # Process the response
+        response = HTTPRequests.post(url, headers=headers, json=data)
         if response.status_code == 200:
-            # Extract embeddings from the response and convert them to a single NumPy array
             embeddings = response.json().get('data', [])
-            embeddings_np = np.array([embedding_data['embedding'] for embedding_data in embeddings], dtype="f")
+            embeddings_np = np.array([e['embedding'] for e in embeddings], dtype='f')
             return embeddings_np
         elif response.status_code == 429:
             raise HTTPRequests.exceptions.RequestException(
                 f"Rate limit exceeded: {response.status_code}, {response.text}"
             )
-        
-        ## When the input is too long, we need to segment the text
         elif response.status_code == 400:
+            # input too long; segment and average
             ebd = []
             for text in texts:
                 chunks = self._Jina_segmenter(text, max_token=MAX_TOKENS)
@@ -140,90 +164,68 @@ class Embedder:
                 chunk_embedding = self.get_embedding(chunks)
                 weighted_avg = np.average(chunk_embedding, weights=token_counts, axis=0)
                 ebd.append(weighted_avg)
-            return np.array(ebd, dtype="f")
-            
+            return np.array(ebd, dtype='f')
         else:
-            print(f"Error: {response.status_code}, {response.text}")
             raise Exception(f"Failed to get embedding from Jina API: {response.status_code}, {response.text}")
-    
+
     @on_exception(expo, HTTPRequests.exceptions.RequestException, max_time=30)
-    def _get_openai_embedding(self, texts: list) -> np.ndarray:
-        """Fetches embeddings from the OpenAI API and returns them as a NumPy array. Requires OpenAI API key in .env file."""
-        # openai embedding API has a limit on single batch size of 2048 texts, so we may need to batch here
+    def _get_openai_embedding(self, texts: List[str]) -> np.ndarray:
+        """Fetch embeddings from the OpenAI API and return them as a NumPy array."""
         batch_size = 2048
         if len(texts) > batch_size:
             all_embeddings = []
             for i in range(0, len(texts), batch_size):
-                batch_texts = texts[i:i+batch_size]
+                batch_texts = texts[i:i + batch_size]
                 batch_embeddings = self._get_openai_embedding(batch_texts)
                 all_embeddings.append(batch_embeddings)
             return np.vstack(all_embeddings)
-        
-        texts = [text.replace("\n", " ") for text in texts]  # Clean text input
+        texts = [text.replace("\n", " ") for text in texts]
         response = self.client.embeddings.create(input=texts, model=self.model_name)
-
-        # Extract embeddings from response
         embeddings = [item.embedding for item in response.data]
+        return np.array(embeddings, dtype='f')
 
-        # Convert the list of embeddings to a NumPy array with the desired data type
-        return np.array(embeddings, dtype="f")
+    def _get_torch_embedding(self, texts: List[str]) -> np.ndarray:
+        """Generate embeddings using a local PyTorch model."""
+        import torch
+        import torch.nn.functional as F
 
-    def _get_torch_embedding(self, texts: list) -> np.ndarray:
-        """Generates embeddings using a local PyTorch model."""
-        import torch  # Importing PyTorch only when needed
         @torch.no_grad()
-        def _encode(self, input_texts):
-            """
-            Generates embeddings for a list of texts using a pytorch local model.
-            
-            Args:
-                input_texts (list of str): A list of texts to encode.
-            
-            Returns:
-                np.ndarray: An array of embeddings.
-            """
-            batch_dict = self.tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt', return_attention_mask=True).to('cuda' if self.to_cuda else 'cpu')
-            
-            outputs = self.model(**batch_dict)
-            attention_mask = batch_dict["attention_mask"]
+        def _encode(instance: 'Embedder', input_texts: List[str]) -> np.ndarray:
+            batch_dict = instance.tokenizer(
+                input_texts,
+                max_length=512,
+                padding=True,
+                truncation=True,
+                return_tensors='pt',
+                return_attention_mask=True,
+            ).to('cuda' if instance.to_cuda else 'cpu')
+            outputs = instance.model(**batch_dict)
+            attention_mask = batch_dict['attention_mask']
             hidden = outputs.last_hidden_state
-
-            reps = _weighted_mean_pooling(hidden, attention_mask)   
+            def _weighted_mean_pooling(hidden_states, mask):
+                # compute weighted mean over tokens
+                mask_ = mask * mask.cumsum(dim=1)
+                s = (hidden_states * mask_.unsqueeze(-1).float()).sum(dim=1)
+                d = mask_.sum(dim=1, keepdim=True).float()
+                return s / d
+            reps = _weighted_mean_pooling(hidden, attention_mask)
             embeddings = F.normalize(reps, p=2, dim=1).detach().cpu().numpy()
             return embeddings
-        
-        def _weighted_mean_pooling(hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-            """
-            Computes weighted mean pooling over the hidden states.
-            
-            Args:
-                hidden (torch.Tensor): The hidden states output from the transformer model.
-                attention_mask (torch.Tensor): The attention mask for the input sequences.
-            
-            Returns:
-                torch.Tensor: The pooled representation of the input.
-            """
-            attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
-            s = torch.sum(hidden * attention_mask_.unsqueeze(-1).float(), dim=1)
-            d = attention_mask_.sum(dim=1, keepdim=True).float()
-            reps = s / d
-            return reps
-        
         return _encode(self, texts)
-    
+
     @on_exception(expo, HTTPRequests.exceptions.RequestException, max_time=30)
-    def _Jina_segmenter(self, text: str, max_token: int) -> list[str]:
-        """Segments text into chunks using Jina API. (free but need API key)"""
+    def _Jina_segmenter(self, text: str, max_token: int) -> List[str]:
+        """Segments text into chunks using Jina API.  (free but needs API key)"""
         url = 'https://segment.jina.ai/'
         headers = {
             'Content-Type': 'application/json',
-            'Authorization': f'Bearer {getenv("JINAAI_API_KEY")}'
+            'Authorization': f"Bearer {getenv('JINAAI_API_KEY')}"
         }
         data = {
             "content": text,
             "return_tokens": True,
             "return_chunks": True,
-            "max_chunk_length": max_token
+            "max_chunk_length": max_token,
         }
         response = HTTPRequests.post(url, headers=headers, json=data)
         return response.json().get('chunks', [])

POP/__init__.py

@@ -1,22 +1,40 @@
-from .POP import PromptFunction, get_text_snapshot
-from .Embedder import Embedder
-from .LLMClient import (
-    LLMClient,
-    OpenAIClient,
-    GeminiClient,
-    DeepseekClient,
-    LocalPyTorchClient,
-    DoubaoClient,
+"""Top‑level package for the restructured POP library.
+
+This package exposes the main classes and helper functions for creating
+prompt functions, embeddings and conversation contexts.  It also
+re‑exports provider registry functions for convenience.
+
+Example usage::
+
+    from pop import PromptFunction, Context, list_providers
+
+    ctx = Context(system="You are a helpful assistant")
+    pf = PromptFunction(sys_prompt="Translate", prompt="<<<text>>>", client="openai")
+    result = pf.execute(text="Hello")
+    print(result)
+"""
+
+from .prompt_function import PromptFunction
+from .embedder import Embedder
+from .context import Context, MessageBlock
+from .api_registry import (
+    list_providers,
+    list_default_model,
+    list_models,
+    get_default_model,
+    get_model,
+    get_client,
 )
 
 __all__ = [
     "PromptFunction",
-    "get_text_snapshot",    
     "Embedder",
-    "LLMClient",
-    "OpenAIClient",
-    "GeminiClient",
-    "DeepseekClient",
-    "LocalPyTorchClient",
-    "DoubaoClient",
+    "Context",
+    "MessageBlock",
+    "list_providers",
+    "list_default_model",
+    "list_models",
+    "get_default_model",
+    "get_model",
+    "get_client",
 ]

POP/api_registry.py

@@ -0,0 +1,148 @@
+"""Provider and model registry.
+
+This module offers helper functions to inspect and instantiate LLM
+providers.  It borrows the concept of a central registry from the
+``api‑registry.ts`` file in the pi‑ai project and exposes:
+
+* :func:`list_providers` – return a list of provider identifiers.
+* :func:`list_default_model` – return a mapping of provider identifiers
+  to their default model names.
+* :func:`_get_default_model` – return the default model for a provider.
+* :func:`list_models` – return a mapping of provider identifiers to
+  all available model names from ``providers.json``.
+* :func:`get_model` – instantiate and return a client based on a model name.
+* :func:`get_client` – instantiate and return a client class given
+  its provider identifier.
+"""
+
+from __future__ import annotations
+
+import json
+from os import path
+from typing import Dict, List, Optional
+
+from .providers import DEFAULT_CLIENTS
+from .models import DEFAULT_MODEL
+
+_PROVIDERS_JSON = path.join(path.dirname(__file__), "providers", "providers.json")
+
+
+def _load_provider_catalog() -> Dict[str, Dict[str, object]]:
+    if not path.exists(_PROVIDERS_JSON):
+        return {}
+    with open(_PROVIDERS_JSON, "r", encoding="utf-8") as handle:
+        data = json.load(handle)
+    if not isinstance(data, dict):
+        return {}
+    return data
+
+
+def _flatten_models(provider_data: Dict[str, object]) -> List[str]:
+    models: List[str] = []
+    seen: set[str] = set()
+    for value in provider_data.values():
+        if isinstance(value, list) and all(isinstance(item, str) for item in value):
+            for item in value:
+                if item not in seen:
+                    seen.add(item)
+                    models.append(item)
+    return models
+
+
+def list_providers() -> List[str]:
+    """Return the list of registered provider identifiers."""
+    return list(DEFAULT_CLIENTS.keys())
+
+
+def list_default_model() -> Dict[str, str]:
+    """Return a mapping of provider identifiers to default model names."""
+    models: Dict[str, str] = {}
+    for name, cls in DEFAULT_CLIENTS.items():
+        model_name = DEFAULT_MODEL.get(cls.__name__)
+        if model_name:
+            models[name] = model_name
+    return models
+
+
+def _get_default_model(provider_name: str) -> Optional[str]:
+    """Return the default model name for a provider."""
+    cls = DEFAULT_CLIENTS.get(provider_name)
+    if cls is None:
+        return None
+    return DEFAULT_MODEL.get(cls.__name__)
+
+
+def list_models() -> Dict[str, List[str]]:
+    """Return a mapping of provider identifiers to available model names."""
+    providers = _load_provider_catalog()
+    models: Dict[str, List[str]] = {}
+    for provider_name, provider_data in providers.items():
+        if isinstance(provider_data, dict):
+            flattened = _flatten_models(provider_data)
+            if flattened:
+                models[provider_name] = flattened
+    return models
+
+def get_default_model(provider_name: str) -> Optional[str]:
+    """Return the default model name for a provider."""
+    return _get_default_model(provider_name)
+
+def get_model(model_name: str) -> Optional[object]:
+    """Instantiate and return a client based on a model name.
+
+    The lookup searches the provider catalog first, then falls back to
+    default model names defined in ``DEFAULT_MODEL``.
+    """
+    # Search providers.json for a matching model name
+    providers = _load_provider_catalog()
+    for provider_name, provider_data in providers.items():
+        if isinstance(provider_data, dict):
+            flattened = _flatten_models(provider_data)
+            if model_name in flattened:
+                return get_client(provider_name, model_name)
+    # Fall back to default model mappings
+    for provider_name, cls in DEFAULT_CLIENTS.items():
+        default_model = DEFAULT_MODEL.get(cls.__name__)
+        if default_model == model_name:
+            return get_client(provider_name, model_name)
+    return None
+
+def get_client(provider_name: str, model_name: Optional[str] = None) -> Optional[object]:
+    """Instantiate and return an LLM client for the given provider.
+
+    Parameters
+    ----------
+    provider_name : str
+        The provider identifier (e.g. ``"openai"``).
+
+    model_name : str
+        The model name (available in provider.json).
+
+    Returns
+    -------
+    object | None
+        An instance of the provider's client class if recognised,
+        otherwise ``None``.
+    """
+    if model_name is None:
+        model_name = _get_default_model(provider_name)  
+    cls = DEFAULT_CLIENTS.get(provider_name)
+
+    if cls is not None:
+        try:
+            return cls(model=model_name)
+        except TypeError:
+            # Some clients (e.g., LocalPyTorchClient) do not accept model args.
+            return cls()
+    
+    return None
+
+
+__all__ = [
+    "list_providers",
+    "list_default_model",
+    "list_models",
+    "get_model",
+    "get_default_model",
+    "get_client",
+]

POP/context.py

@@ -0,0 +1,47 @@
+"""Conversation context objects.
+
+In order to build and track a conversation with an LLM, this module
+defines simple data classes for system messages and user/assistant
+messages.  A :class:`Context` holds a list of :class:`MessageBlock`
+objects and optionally a system prompt and a list of tool
+descriptions.  It can be converted into the message list expected by
+``LLMClient.chat_completion``.
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any
+
+
+@dataclass
+class MessageBlock:
+    """Represents a single message from a user or assistant."""
+    role: str
+    content: str
+
+
+@dataclass
+class Context:
+    """Collects conversation messages and metadata."""
+    system: Optional[str] = None
+    messages: List[MessageBlock] = field(default_factory=list)
+    tools: Optional[List[Dict[str, Any]]] = None
+
+    def append(self, role: str, content: str) -> None:
+        """Append a message to the context."""
+        self.messages.append(MessageBlock(role=role, content=content))
+
+    def to_messages(self) -> List[Dict[str, Any]]:
+        """Convert the context into a list of message dictionaries.
+
+        The resulting list begins with a system message if one is set,
+        followed by all appended messages in order.
+        """
+        output: List[Dict[str, Any]] = []
+        if self.system:
+            output.append({"role": "system", "content": self.system})
+        for mb in self.messages:
+            output.append({"role": mb.role, "content": mb.content})
+        return output
+
+
+__all__ = ["MessageBlock", "Context"]

POP/env_api_keys.py

@@ -0,0 +1,33 @@
+"""Environment key helpers.
+
+Providers may require API keys to function.  This module defines
+utility functions to inspect whether the necessary keys are present in
+the environment.  Consumers can call :func:`has_api_key` to check
+whether a provider is ready for use.
+"""
+
+import os
+from typing import Optional
+from dotenv import load_dotenv
+
+if not load_dotenv():
+    print("No .env file found or could not be loaded.")
+
+# Mapping from provider identifier to the environment variable used
+REQUIRED_KEYS = {
+    "openai": "OPENAI_API_KEY",
+    "gemini": "GEMINI_API_KEY",
+    "deepseek": "DEEPSEEK_API_KEY",
+    "doubao": "DOUBAO_API_KEY",
+    # local and ollama do not require API keys by default
+}
+
+
+
+def has_api_key(provider: str) -> bool:
+    """Return True if the required API key for *provider* is set."""
+    env_var = REQUIRED_KEYS.get(provider)
+    return bool(os.getenv(env_var)) if env_var else True
+
+
+__all__ = ["has_api_key"]

POP/models.py

@@ -0,0 +1,20 @@
+"""Model names and default values.
+
+This module centralises the default model names used for each
+provider.  It mirrors the ``default_model`` dictionary from the
+original POP implementation.  Other modules import this to obtain a
+provider's default model.
+"""
+
+# Map provider class names to their default model identifiers
+DEFAULT_MODEL = {
+    "OpenAIClient": "gpt-5-nano",
+    "GeminiClient": "gemini-2.5-flash",
+    "DeepseekClient": "deepseek-chat",
+    "DoubaoClient": "doubao-seed-1-6-flash-250715",
+    "OllamaClient": "mistral:7b",
+    # LocalPyTorchClient is a stub; choose an arbitrary default
+    "LocalPyTorchClient": "local-llm",
+}
+
+__all__ = ["DEFAULT_MODEL"]