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

parrot/stores/utils/chunking.py

@@ -0,0 +1,197 @@
+from typing import List, Dict, Any, Tuple, Optional
+from dataclasses import dataclass
+import re
+import numpy as np
+
+
+@dataclass
+class ChunkInfo:
+    """Information about a document chunk"""
+    chunk_id: str
+    parent_document_id: str
+    chunk_index: int
+    chunk_text: str
+    start_position: int
+    end_position: int
+    chunk_embedding: np.ndarray
+    metadata: Dict[str, Any]
+
+
+class LateChunkingProcessor:
+    """
+    Late Chunking processor integrated with PgVectorStore.
+
+    Late chunking generates embeddings for the full document first, then creates
+    contextually-aware chunk embeddings that preserve the global document context.
+    """
+
+    def __init__(
+        self,
+        vector_store,
+        chunk_size: int = 8192,
+        chunk_overlap: int = 200,
+        preserve_sentences: bool = True,
+        min_chunk_size: int = 100
+    ):
+        self.vector_store = vector_store
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.preserve_sentences = preserve_sentences
+        self.min_chunk_size = min_chunk_size
+
+    async def process_document_late_chunking(
+        self,
+        document_text: str,
+        document_id: str,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> Tuple[np.ndarray, List[ChunkInfo]]:
+        """
+        Process document with late chunking strategy.
+
+        Args:
+            document_text: Full document text
+            document_id: Unique document identifier
+            metadata: Optional metadata for the document
+
+        Returns:
+            Tuple of (full_document_embedding, list_of_chunk_info)
+        """
+        # Step 1: Generate full-document embedding for global context
+        full_embedding = self.vector_store._embed_.embed_query(document_text)
+
+        # Step 2: Split into semantic chunks
+        chunks = self._semantic_chunk_split(document_text)
+
+        # Step 3: Generate contextual embeddings for each chunk
+        chunk_infos = []
+
+        for chunk_idx, (chunk_text, start_pos, end_pos) in enumerate(chunks):
+            # Create contextual prompt that includes document context
+            contextual_text = self._create_contextual_text(
+                document_text, chunk_text, start_pos, end_pos
+            )
+
+            # Generate embedding with context
+            chunk_embedding = self.vector_store._embed_.embed_query(contextual_text)
+
+            # Create chunk ID
+            chunk_id = f"{document_id}_chunk_{chunk_idx:04d}"
+
+            # Prepare chunk metadata
+            chunk_metadata = {
+                **(metadata or {}),
+                'parent_document_id': document_id,
+                'chunk_index': chunk_idx,
+                'total_chunks': len(chunks),
+                'start_position': start_pos,
+                'end_position': end_pos,
+                'chunk_size': len(chunk_text),
+                'is_chunk': True,
+                'chunk_type': 'late_chunking',
+                'context_preserved': True
+            }
+
+            chunk_info = ChunkInfo(
+                chunk_id=chunk_id,
+                parent_document_id=document_id,
+                chunk_index=chunk_idx,
+                chunk_text=chunk_text,
+                start_position=start_pos,
+                end_position=end_pos,
+                chunk_embedding=chunk_embedding,
+                metadata=chunk_metadata
+            )
+
+            chunk_infos.append(chunk_info)
+
+        return np.array(full_embedding), chunk_infos
+
+    def _semantic_chunk_split(self, text: str) -> List[Tuple[str, int, int]]:
+        """
+        Split text preserving semantic boundaries.
+
+        Returns:
+            List of (chunk_text, start_position, end_position) tuples
+        """
+        if self.preserve_sentences:
+            return self._sentence_aware_chunking(text)
+        else:
+            return self._simple_chunking(text)
+
+    def _sentence_aware_chunking(self, text: str) -> List[Tuple[str, int, int]]:
+        """Split text while preserving sentence boundaries."""
+        # Split by sentences (basic approach - could use spaCy for better results)
+        sentence_endings = re.finditer(r'[.!?]+\s+', text)
+        sentence_positions = [0] + [m.end() for m in sentence_endings] + [len(text)]
+
+        chunks = []
+        current_start = 0
+
+        for i in range(1, len(sentence_positions)):
+            current_end = sentence_positions[i]
+            current_size = current_end - current_start
+
+            # If current chunk is too large, create chunk at previous boundary
+            if current_size > self.chunk_size and len(chunks) > 0:
+                # Find the last good break point
+                prev_end = sentence_positions[i-1]
+                if prev_end - current_start >= self.min_chunk_size:
+                    chunk_text = text[current_start:prev_end].strip()
+                    chunks.append((chunk_text, current_start, prev_end))
+
+                    # Start new chunk with overlap
+                    overlap_start = max(current_start, prev_end - self.chunk_overlap)
+                    current_start = overlap_start
+
+            # If we're at the end, add final chunk
+            if i == len(sentence_positions) - 1:
+                chunk_text = text[current_start:current_end].strip()
+                if len(chunk_text) >= self.min_chunk_size:
+                    chunks.append((chunk_text, current_start, current_end))
+
+        return chunks if chunks else [(text, 0, len(text))]
+
+    def _simple_chunking(self, text: str) -> List[Tuple[str, int, int]]:
+        """Simple character-based chunking with overlap."""
+        chunks = []
+        start = 0
+
+        while start < len(text):
+            end = min(start + self.chunk_size, len(text))
+            chunk_text = text[start:end].strip()
+
+            if len(chunk_text) >= self.min_chunk_size:
+                chunks.append((chunk_text, start, end))
+
+            # Move start position with overlap
+            start += self.chunk_size - self.chunk_overlap
+
+            if end == len(text):
+                break
+
+        return chunks
+
+    def _create_contextual_text(
+        self,
+        full_text: str,
+        chunk_text: str,
+        start_pos: int,
+        end_pos: int
+    ) -> str:
+        """
+        Create contextual text that includes surrounding context for better embeddings.
+        """
+        # Get surrounding context (e.g., 200 chars before and after)
+        context_window = 200
+
+        context_start = max(0, start_pos - context_window)
+        context_end = min(len(full_text), end_pos + context_window)
+
+        # Extract context
+        before_context = full_text[context_start:start_pos] if context_start < start_pos else ""
+        after_context = full_text[end_pos:context_end] if end_pos < context_end else ""
+
+        # Create contextual text with clear boundaries
+        contextual_text = f"{before_context.strip()} [FOCUS] {chunk_text} [/FOCUS] {after_context.strip()}"  # noqa
+
+        return contextual_text.strip()

parrot/telemetry/__init__.py

@@ -0,0 +1,3 @@
+from .mixin import TelemetryMixin
+
+__all__ = ["TelemetryMixin"]

parrot/telemetry/mixin.py

@@ -0,0 +1,111 @@
+# parrot/telemetry/__init__.py
+from typing import Optional, Dict, Any
+from contextlib import contextmanager
+import time
+import openlit
+from opentelemetry import trace, metrics
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.metrics import MeterProvider
+
+
+class TelemetryMixin:
+    """Mixin to add observability to LLM clients"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._telemetry_enabled = kwargs.get('enable_telemetry', True)
+
+        if self._telemetry_enabled:
+            # Get tracer and meter for this client
+            self.tracer = trace.get_tracer(f"parrot.client.{self.client_name}")
+            self.meter = metrics.get_meter(f"parrot.client.{self.client_name}")
+
+            # Create metrics
+            self.request_counter = self.meter.create_counter(
+                name="llm.requests",
+                description="Total LLM requests",
+                unit="1"
+            )
+            self.token_counter = self.meter.create_counter(
+                name="llm.tokens",
+                description="Total tokens used",
+                unit="1"
+            )
+            self.latency_histogram = self.meter.create_histogram(
+                name="llm.request.duration",
+                description="LLM request duration",
+                unit="ms"
+            )
+            self.error_counter = self.meter.create_counter(
+                name="llm.errors",
+                description="LLM errors",
+                unit="1"
+            )
+
+    @contextmanager
+    def track_request(self, operation: str, **attributes):
+        """Context manager for tracking LLM requests"""
+        if not self._telemetry_enabled:
+            yield {}
+            return
+
+        start_time = time.time()
+        span_attributes = {
+            "llm.provider": self.client_name,
+            "llm.operation": operation,
+            **attributes
+        }
+
+        with self.tracer.start_as_current_span(
+            f"{self.client_name}.{operation}",
+            attributes=span_attributes
+        ) as span:
+            metrics = {
+                "start_time": start_time,
+                "span": span
+            }
+
+            try:
+                yield metrics
+                # Success tracking
+                self.request_counter.add(1, {
+                    "provider": self.client_name,
+                    "operation": operation,
+                    "status": "success"
+                })
+            except Exception as e:
+                # Error tracking
+                self.error_counter.add(1, {
+                    "provider": self.client_name,
+                    "operation": operation,
+                    "error_type": type(e).__name__
+                })
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+            finally:
+                # Always track latency
+                duration_ms = (time.time() - start_time) * 1000
+                self.latency_histogram.record(duration_ms, {
+                    "provider": self.client_name,
+                    "operation": operation
+                })
+
+    def track_tokens(self, usage: Dict[str, int], model: str):
+        """Track token usage"""
+        if not self._telemetry_enabled:
+            return
+
+        if "input_tokens" in usage:
+            self.token_counter.add(usage["input_tokens"], {
+                "provider": self.client_name,
+                "model": model,
+                "token_type": "input"
+            })
+
+        if "output_tokens" in usage:
+            self.token_counter.add(usage["output_tokens"], {
+                "provider": self.client_name,
+                "model": model,
+                "token_type": "output"
+            })

parrot/template/__init__.py

@@ -0,0 +1,3 @@
+from .engine import TemplateEngine, JinjaConfig
+
+__all__ = ["TemplateEngine", "JinjaConfig"]

parrot/template/engine.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Union
+import datetime
+from navconfig.logging import logging
+
+from jinja2 import (
+    BaseLoader,
+    ChoiceLoader,
+    DictLoader,
+    Environment,
+    FileSystemBytecodeCache,
+    FileSystemLoader,
+    TemplateError,
+    TemplateNotFound,
+    StrictUndefined,
+    select_autoescape,
+)
+
+PathLike = Union[str, Path]
+
+
+@dataclass
+class JinjaConfig:
+    """Configuration for the async Jinja2 Environment."""
+    template_dirs: list[Path] = field(default_factory=list)
+    extensions: list[str] = field(default_factory=lambda: [
+        "jinja2.ext.i18n",
+        "jinja2.ext.loopcontrols",
+        "jinja2_time.TimeExtension",
+        "jinja2_iso8601.ISO8601Extension",
+        "jinja2.ext.do",
+        "jinja2_humanize_extension.HumanizeExtension",
+        # "jinja2.ext.debug",  # enable in dev if desired
+    ])
+    bytecode_cache_dir: Optional[Path] = None
+    bytecode_cache_pattern: str = "%s.cache"
+    autoescape: Any = select_autoescape(["html", "xml", "j2", "jinja", "jinja2"])
+    undefined: Any = StrictUndefined  # raise on missing variables
+    keep_trailing_newline: bool = True
+    trim_blocks: bool = True
+    lstrip_blocks: bool = True
+
+
+class TemplateEngine:
+    """
+    Async-only Jinja2 template engine with:
+        - multiple directories
+        - in-memory templates
+        - pluggable extensions/filters/globals
+        - optional bytecode cache
+    """
+
+    def __init__(
+        self,
+        template_dirs: Optional[Union[PathLike, Sequence[PathLike]]] = None,
+        *,
+        extensions: Optional[Sequence[str]] = None,
+        bytecode_cache_dir: Optional[PathLike] = None,
+        filters: Optional[Mapping[str, Any]] = None,
+        globals_: Optional[Mapping[str, Any]] = None,
+        config: Optional[JinjaConfig] = None,
+        debug: bool = False,
+    ) -> None:
+        cfg = config or JinjaConfig()
+        self.logger = logging.getLogger(__name__)
+
+        # Normalize directories
+        if template_dirs is not None:
+            if isinstance(template_dirs, (str, Path)):
+                cfg.template_dirs.append(Path(template_dirs).resolve())
+            else:
+                cfg.template_dirs.extend(Path(p).resolve() for p in template_dirs)
+
+        # Validate & store directories
+        self._fs_dirs: list[Path] = []
+        for d in cfg.template_dirs:
+            if not d.exists():
+                raise ValueError(f"Template directory not found: {d}")
+            if not d.is_dir():
+                raise ValueError(f"Template path is not a directory: {d}")
+            self._fs_dirs.append(d)
+
+        # Extensions
+        self._configure_extensions(cfg, extensions, debug)
+
+        # Optional bytecode cache
+        self._bytecode_cache = None
+        if bytecode_cache_dir is not None:
+            bdir = Path(bytecode_cache_dir).resolve()
+            bdir.mkdir(parents=True, exist_ok=True)
+            self._bytecode_cache = FileSystemBytecodeCache(str(bdir), cfg.bytecode_cache_pattern)
+        elif cfg.bytecode_cache_dir:
+            bdir = Path(cfg.bytecode_cache_dir).resolve()
+            bdir.mkdir(parents=True, exist_ok=True)
+            self._bytecode_cache = FileSystemBytecodeCache(str(bdir), cfg.bytecode_cache_pattern)
+
+        # Loaders: FileSystem + (optional) in-memory DictLoader
+        self._dict_loader = DictLoader({})
+        fs_loader = FileSystemLoader([str(p) for p in self._fs_dirs]) if self._fs_dirs else None
+
+        if fs_loader:
+            loader: BaseLoader = ChoiceLoader([self._dict_loader, fs_loader])
+        else:
+            loader = self._dict_loader  # still works with only in-memory templates
+
+        # Build the environment (async-only)
+        self.env = Environment(
+            loader=loader,
+            enable_async=True,
+            extensions=self._extensions,
+            bytecode_cache=self._bytecode_cache,
+            autoescape=cfg.autoescape,
+            undefined=cfg.undefined,
+            keep_trailing_newline=cfg.keep_trailing_newline,
+            trim_blocks=cfg.trim_blocks,
+            lstrip_blocks=cfg.lstrip_blocks,
+        )
+
+        # Useful default filters/globals
+        self.env.filters.setdefault("datetime", datetime.datetime.fromtimestamp)
+        # Add user-supplied
+        if filters:
+            for name, fn in filters.items():
+                self.env.filters[name] = fn
+        if globals_:
+            for name, val in globals_.items():
+                self.env.globals[name] = val
+
+        self.logger.debug(
+            "AsyncTemplateEngine initialized: dirs=%s, extensions=%s, bytecode_cache=%s",
+            [str(d) for d in self._fs_dirs],
+            self._extensions,
+            bool(self._bytecode_cache),
+        )
+
+    def _configure_extensions(self, cfg: JinjaConfig, extensions: Optional[Sequence[str]], debug: bool) -> None:
+        """Configure Jinja2 extensions."""
+        # Merge extensions
+        self._extensions: list[str] = list(cfg.extensions)
+        if extensions:
+            for ext in extensions:
+                if ext not in self._extensions:
+                    self._extensions.append(ext)
+        if debug and "jinja2.ext.debug" not in self._extensions:
+            self._extensions.append("jinja2.ext.debug")
+
+    # -------- Public API
+    def add_template_dir(self, path: PathLike) -> None:
+        """Add a new filesystem directory to the search path at runtime."""
+        p = Path(path).resolve()
+        if not p.exists() or not p.is_dir():
+            raise ValueError(f"Template directory invalid: {p}")
+        self._fs_dirs.append(p)
+
+        # Rebuild loader chain (ChoiceLoader is immutable in practice)
+        current_dicts = self._dict_loader.mapping
+        self._dict_loader = DictLoader(current_dicts)
+        fs_loader = FileSystemLoader([str(d) for d in self._fs_dirs])
+        self.env.loader = ChoiceLoader([self._dict_loader, fs_loader])
+
+    def add_templates(self, templates: Mapping[str, str]) -> None:
+        """
+        Add/override in-memory templates.
+        Example: add_templates({'layout.html': '<html>{{ self.block() }}</html>'})
+        """
+        for name, content in templates.items():
+            self._dict_loader.mapping[name] = content
+
+    def get_template(self, name: str):
+        """Get a compiled template by name (raises FileNotFoundError on miss)."""
+        try:
+            return self.env.get_template(name)
+        except TemplateNotFound as ex:
+            raise FileNotFoundError(f"Template not found: {name}") from ex
+        except Exception as ex:
+            raise RuntimeError(f"Error loading template '{name}': {ex}") from ex
+
+    async def render(self, name: str, params: Optional[Mapping[str, Any]] = None) -> str:
+        """
+        Async render of a template by name.
+        Only async path is supported (no sync render).
+        """
+        params = dict(params or {})
+        try:
+            tmpl = self.get_template(name)
+            # MUST use render_async because enable_async=True
+            return await tmpl.render_async(**params)
+        except TemplateError as ex:
+            raise ValueError(f"Template error while rendering '{name}': {ex}") from ex
+        except Exception as ex:
+            raise RuntimeError(f"Unexpected error rendering '{name}': {ex}") from ex
+
+    async def render_string(self, source: str, params: Optional[Mapping[str, Any]] = None) -> str:
+        """
+        Async render from a string (compiled via the current environment).
+        Useful for ad-hoc/injected content.
+        """
+        params = dict(params or {})
+        try:
+            tmpl = self.env.from_string(source)
+            return await tmpl.render_async(**params)
+        except TemplateError as ex:
+            raise ValueError(f"Template error while rendering string: {ex}") from ex
+        except Exception as ex:
+            raise RuntimeError(f"Unexpected error rendering string: {ex}") from ex
+
+    def add_filters(self, filters: Mapping[str, Any]) -> None:
+        """Register additional filters (supports async filters too)."""
+        for name, fn in filters.items():
+            self.env.filters[name] = fn
+
+    def add_globals(self, globals_: Mapping[str, Any]) -> None:
+        """Register additional global variables/functions."""
+        for name, val in globals_.items():
+            self.env.globals[name] = val
+
+    def compile_directory(self, target: PathLike, *, zip: Optional[str] = "deflated") -> None:
+        """
+        Optionally precompile all templates from filesystem loaders into `target`.
+        Skips silently if there are no filesystem directories.
+        """
+        if not self._fs_dirs:
+            return
+        target_path = Path(target).resolve()
+        target_path.mkdir(parents=True, exist_ok=True)
+        try:
+            self.env.compile_templates(
+                target=str(target_path),
+                zip=zip
+            )
+        except Exception as ex:
+            # If a transient decode error happens, try once again without failing the app
+            self.logger.warning(
+                f"Compile templates failed once: {ex}; retrying…"
+            )
+            self.env.compile_templates(
+                target=str(target_path),
+                zip=zip
+            )
+
+
+# ------------------------------
+# Minimal usage example (async)
+# ------------------------------
+# engine = TemplateEngine(
+#     template_dirs=["/path/to/templates", "/path/to/partials"],
+#     extensions=["jinja2.ext.debug"],  # optional
+#     bytecode_cache_dir="/tmp/jinja-cache",  # optional
+#     filters={"jsonify": json_encoder},      # e.g. your custom filter
+#     globals_={"app_name": "QuerySource"},   # optional
+# )
+#
+# engine.add_templates({"inline.html": "Hello {{ name }}!"})
+#
+# html = await engine.render("index.html", {"user": "Pilar"})
+# text = await engine.render_string("Hi {{ who }}", {"who": "there"})

parrot/tools/__init__.py

@@ -0,0 +1,23 @@
+"""
+Tools infrastructure for building Agents.
+"""
+from parrot.plugins import setup_plugin_importer, dynamic_import_helper
+from .abstract import AbstractTool, ToolResult
+from .toolkit import AbstractToolkit, ToolkitTool
+from .decorators import tool_schema, tool
+
+setup_plugin_importer('parrot.tools', 'tools')
+
+__all__ = (
+    "AbstractTool",
+    "ToolResult",
+    "AbstractToolkit",
+    "ToolkitTool",
+    "tool_schema",
+    "tool",
+)
+
+# Enable dynamic imports
+def __getattr__(name):
+    return dynamic_import_helper(__name__, name)
+