proscenium 0.0.8__tar.gz → 0.0.9__tar.gz

{proscenium-0.0.8 → proscenium-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: proscenium
-Version: 0.0.8
+Version: 0.0.9
 Summary: Frame AI Agents
 License: ASFv2
 Author: Adam Pingel
@@ -18,6 +18,7 @@ Requires-Dist: docstring_parser (>=0.16,<0.17)
 Requires-Dist: gofannon (>=0.25.13,<0.26.0)
 Requires-Dist: langchain-community (>=0.3.18,<0.4.0)
 Requires-Dist: langchain-huggingface (>=0.1.2,<0.2.0)
+Requires-Dist: lapidarist (>=0.0.1,<0.0.2)
 Requires-Dist: neo4j (>=5.28.1,<6.0.0)
 Requires-Dist: pydantic (>=2.10.6,<3.0.0)
 Requires-Dist: pymilvus (>=2.5.4,<3.0.0)

{proscenium-0.0.8 → proscenium-0.0.9}/proscenium/patterns/rag.py

@@ -4,9 +4,10 @@ import logging
 from pymilvus import MilvusClient
 from pymilvus import model
 
+from lapidarist.verbs.display.milvus import chunk_hits_table
+from lapidarist.verbs.vector_database import closest_chunks
+
 from proscenium.verbs.complete import complete_simple
-from proscenium.verbs.display.milvus import chunk_hits_table
-from proscenium.verbs.vector_database import closest_chunks
 
 log = logging.getLogger(__name__)
 

{proscenium-0.0.8 → proscenium-0.0.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "proscenium"
-version = "0.0.8"
+version = "0.0.9"
 description = "Frame AI Agents"
 authors = ["Adam Pingel <oss@pingel.org>"]
 license = "ASFv2"
@@ -9,6 +9,7 @@ packages = [{include = "proscenium"}]
 
 [tool.poetry.dependencies]
 python = "^3.11"
+lapidarist = "^0.0.1"
 aisuite = {extras = ["openai,anthropic"], version = "^0.1.10"}
 rich = "^13.9.4"
 typer = "^0.15.2"

proscenium-0.0.8/proscenium/patterns/chunk_space.py

@@ -1,51 +0,0 @@
-from typing import Optional
-import logging
-from rich.console import Console
-from pymilvus import model
-
-from proscenium.verbs.read import load_file
-from proscenium.verbs.chunk import documents_to_chunks_by_characters
-from proscenium.verbs.display.milvus import collection_panel
-from proscenium.verbs.vector_database import vector_db
-from proscenium.verbs.vector_database import create_collection
-from proscenium.verbs.vector_database import add_chunks_to_vector_db
-
-log = logging.getLogger(__name__)
-
-
-def load_chunks_from_files(
-    data_files: list[str],
-    milvus_uri: str,
-    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
-    collection_name: str,
-    console: Optional[Console] = None,
-) -> None:
-
-    vector_db_client = vector_db(milvus_uri)
-    log.info("Vector db stored at %s", milvus_uri)
-
-    for data_file in data_files:
-
-        log.info(
-            "Loading data file %s into vector db %s collection %s",
-            data_file,
-            milvus_uri,
-            collection_name,
-        )
-        create_collection(vector_db_client, embedding_fn, collection_name)
-
-        documents = load_file(data_file)
-        chunks = documents_to_chunks_by_characters(documents)
-        log.info("Data file %s has %s chunks", data_file, len(chunks))
-
-        info = add_chunks_to_vector_db(
-            vector_db_client,
-            embedding_fn,
-            chunks,
-            collection_name,
-        )
-        log.info("%s chunks inserted ", info["insert_count"])
-        if console is not None:
-            console.print(collection_panel(vector_db_client, collection_name))
-
-    vector_db_client.close()

proscenium-0.0.8/proscenium/patterns/document_enricher.py

@@ -1,84 +0,0 @@
-from typing import List
-from typing import Callable
-from typing import Optional
-
-import time
-import logging
-from pydantic import BaseModel
-
-from rich.panel import Panel
-from rich.console import Console
-from rich.progress import Progress
-
-from langchain_core.documents.base import Document
-
-from proscenium.verbs.chunk import documents_to_chunks_by_tokens
-from proscenium.verbs.extract import extract_to_pydantic_model
-
-log = logging.getLogger(__name__)
-
-
-def extract_from_document_chunks(
-    doc: Document,
-    doc_as_rich: Callable[[Document], Panel],
-    chunk_extraction_model_id: str,
-    chunk_extraction_template: str,
-    chunk_extract_clazz: type[BaseModel],
-    delay: float,
-    console: Optional[Console] = None,
-) -> List[BaseModel]:
-
-    if console is not None:
-        console.print(doc_as_rich(doc))
-        console.print()
-
-    extract_models = []
-
-    chunks = documents_to_chunks_by_tokens([doc], chunk_size=1000, chunk_overlap=0)
-    for i, chunk in enumerate(chunks):
-
-        ce = extract_to_pydantic_model(
-            chunk_extraction_model_id,
-            chunk_extraction_template,
-            chunk_extract_clazz,
-            chunk.page_content,
-        )
-
-        log.info("Extract model in chunk %s of %s", i + 1, len(chunks))
-        if console is not None:
-            console.print(Panel(str(ce)))
-
-        extract_models.append(ce)
-        time.sleep(delay)
-
-    return extract_models
-
-
-def enrich_documents(
-    retrieve_documents: Callable[[], List[Document]],
-    extract_from_doc_chunks: Callable[[Document], List[BaseModel]],
-    doc_enrichments: Callable[[Document, list[BaseModel]], BaseModel],
-    enrichments_jsonl_file: str,
-    console: Optional[Console] = None,
-) -> None:
-
-    docs = retrieve_documents()
-
-    with Progress() as progress:
-
-        task_enrich = progress.add_task(
-            "[green]Enriching documents...", total=len(docs)
-        )
-
-        with open(enrichments_jsonl_file, "wt") as f:
-
-            for doc in docs:
-
-                chunk_extract_models = extract_from_doc_chunks(doc)
-                enrichments = doc_enrichments(doc, chunk_extract_models)
-                enrichments_json = enrichments.model_dump_json()
-                f.write(enrichments_json + "\n")
-
-                progress.update(task_enrich, advance=1)
-
-        log.info("Wrote document enrichments to %s", enrichments_jsonl_file)

proscenium-0.0.8/proscenium/patterns/entity_resolver.py

@@ -1,95 +0,0 @@
-from typing import Optional
-import logging
-
-from rich.console import Console
-from langchain_core.documents.base import Document
-from neo4j import Driver
-
-from pymilvus import MilvusClient
-
-from proscenium.verbs.vector_database import vector_db
-from proscenium.verbs.vector_database import create_collection
-from proscenium.verbs.vector_database import closest_chunks
-from proscenium.verbs.vector_database import add_chunks_to_vector_db
-from proscenium.verbs.vector_database import embedding_function
-from proscenium.verbs.display.milvus import collection_panel
-
-log = logging.getLogger(__name__)
-
-
-class Resolver:
-
-    def __init__(
-        self,
-        cypher: str,
-        field_name: str,
-        collection_name: str,
-    ):
-        self.cypher = cypher
-        self.field_name = field_name
-        self.collection_name = collection_name
-
-
-def load_entity_resolver(
-    driver: Driver,
-    resolvers: list[Resolver],
-    embedding_model_id: str,
-    milvus_uri: str,
-    console: Optional[Console] = None,
-) -> None:
-
-    vector_db_client = vector_db(milvus_uri)
-    log.info("Vector db stored at %s", milvus_uri)
-
-    embedding_fn = embedding_function(embedding_model_id)
-    log.info("Embedding model %s", embedding_model_id)
-
-    for resolver in resolvers:
-
-        values = []
-        with driver.session() as session:
-            result = session.run(resolver.cypher)
-            new_values = [Document(record[resolver.field_name]) for record in result]
-            values.extend(new_values)
-
-        log.info("Loading entity resolver into vector db %s", resolver.collection_name)
-        create_collection(vector_db_client, embedding_fn, resolver.collection_name)
-
-        info = add_chunks_to_vector_db(
-            vector_db_client, embedding_fn, values, resolver.collection_name
-        )
-        log.info("%s chunks inserted ", info["insert_count"])
-
-        if console is not None:
-            console.print(collection_panel(vector_db_client, resolver.collection_name))
-
-    vector_db_client.close()
-
-
-def find_matching_objects(
-    vector_db_client: MilvusClient,
-    approximate: str,
-    resolver: Resolver,
-) -> Optional[str]:
-
-    log.info("Loading collection", resolver.collection_name)
-    vector_db_client.load_collection(resolver.collection_name)
-
-    log.info(
-        "Finding entity matches for", approximate, "using", resolver.collection_name
-    )
-
-    hits = closest_chunks(
-        vector_db_client,
-        resolver.embedding_fn,
-        approximate,
-        resolver.collection_name,
-        k=5,
-    )
-    # TODO apply distance threshold
-    for match in [head["entity"]["text"] for head in hits[:1]]:
-        log.info("Closest match:", match)
-        return match
-
-    log.info("No match found")
-    return None

proscenium-0.0.8/proscenium/patterns/knowledge_graph.py

@@ -1,41 +0,0 @@
-from typing import Callable
-from typing import Any
-
-import logging
-import json
-from pydantic import BaseModel
-
-from rich.progress import Progress
-
-from neo4j import Driver
-
-log = logging.getLogger(__name__)
-
-
-def load_knowledge_graph(
-    driver: Driver,
-    enrichments_jsonl_file: str,
-    enrichments_clazz: type[BaseModel],
-    doc_enrichments_to_graph: Callable[[Any, BaseModel], None],
-) -> None:
-
-    log.info("Parsing enrichments from %s", enrichments_jsonl_file)
-
-    enrichmentss = []
-    with open(enrichments_jsonl_file, "r") as f:
-        for line in f:
-            e = enrichments_clazz.model_construct(**json.loads(line))
-            enrichmentss.append(e)
-
-    with Progress() as progress:
-
-        task_load = progress.add_task(
-            f"Loading {len(enrichmentss)} enriched documents into graph...",
-            total=len(enrichmentss),
-        )
-
-        with driver.session() as session:
-            session.run("MATCH (n) DETACH DELETE n")  # empty graph
-            for e in enrichmentss:
-                session.execute_write(doc_enrichments_to_graph, e)
-                progress.update(task_load, advance=1)

proscenium-0.0.8/proscenium/verbs/chunk.py

@@ -1,42 +0,0 @@
-import logging
-import os
-from typing import List
-from typing import Iterable
-
-from langchain_core.documents.base import Document
-
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.text_splitter import TokenTextSplitter
-
-log = logging.getLogger(__name__)
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-logging.getLogger("langchain_text_splitters.base").setLevel(logging.ERROR)
-
-# Each text chunk inherits the metadata from the document.
-
-
-def documents_to_chunks_by_characters(
-    documents: Iterable[Document], chunk_size: int = 1000, chunk_overlap: int = 0
-) -> List[Document]:
-
-    text_splitter = CharacterTextSplitter(
-        chunk_size=chunk_size, chunk_overlap=chunk_overlap
-    )
-
-    chunks = text_splitter.split_documents(documents)
-
-    return chunks
-
-
-def documents_to_chunks_by_tokens(
-    documents: Iterable[Document], chunk_size: int = 1000, chunk_overlap: int = 0
-) -> List[Document]:
-
-    text_splitter = TokenTextSplitter(
-        chunk_size=chunk_size, chunk_overlap=chunk_overlap
-    )
-
-    chunks = text_splitter.split_documents(documents)
-
-    return chunks

proscenium-0.0.8/proscenium/verbs/display/milvus.py

@@ -1,68 +0,0 @@
-from rich.table import Table
-from rich.panel import Panel
-from rich.text import Text
-from rich.console import Group
-from pymilvus import MilvusClient
-
-
-def chunk_hits_table(chunks: list[dict]) -> Table:
-
-    table = Table(title="Closest Chunks", show_lines=True)
-    table.add_column("id", justify="right")
-    table.add_column("distance")
-    table.add_column("entity.text", justify="right")
-    for chunk in chunks:
-        table.add_row(str(chunk["id"]), str(chunk["distance"]), chunk["entity"]["text"])
-    return table
-
-
-def collection_panel(client: MilvusClient, collection_name: str) -> Panel:
-
-    stats = client.get_collection_stats(collection_name)
-    desc = client.describe_collection(collection_name)
-
-    params_text = Text(
-        f"""
-    Collection Name: {desc['collection_name']}
-    Auto ID: {desc['auto_id']}
-    Num Shards: {desc['num_shards']}
-    Description: {desc['description']}
-    Functions: {desc['functions']}
-    Aliases: {desc['aliases']}
-    Collection ID: {desc['collection_id']}
-    Consistency Level: {desc['consistency_level']}
-    Properties: {desc['properties']}
-    Num Partitions: {desc['num_partitions']}
-    Enable Dynamic Field: {desc['enable_dynamic_field']}"""
-    )
-
-    params_panel = Panel(params_text, title="Params")
-
-    fields_table = Table(title="Fields", show_lines=True)
-    fields_table.add_column("id", justify="left")
-    fields_table.add_column("name", justify="left")
-    fields_table.add_column("description", justify="left")
-    fields_table.add_column("type", justify="left")
-    fields_table.add_column("params", justify="left")
-    fields_table.add_column("auto_id", justify="left")
-    fields_table.add_column("is_primary", justify="left")
-    for field in desc["fields"]:
-        fields_table.add_row(
-            str(field["field_id"]),  # int
-            field["name"],
-            field["description"],
-            field["type"].name,  # Milvus DataType
-            "\n".join([f"{k}: {v}" for k, v in field["params"].items()]),
-            str(field.get("auto_id", "-")),  # bool
-            str(field.get("is_primary", "-")),
-        )  # bool
-
-    stats_text = Text("\n".join([f"{k}: {v}" for k, v in stats.items()]))
-    stats_panel = Panel(stats_text, title="Stats")
-
-    panel = Panel(
-        Group(params_panel, fields_table, stats_panel),
-        title=f"Collection {collection_name}",
-    )
-
-    return panel

proscenium-0.0.8/proscenium/verbs/display/neo4j.py

@@ -1,25 +0,0 @@
-from typing import List
-from rich.table import Table
-
-
-def triples_table(triples: List[tuple[str, str, str]], title: str) -> Table:
-
-    table = Table(title=title, show_lines=False)
-    table.add_column("Subject", justify="left")
-    table.add_column("Predicate", justify="left")
-    table.add_column("Object", justify="left")
-    for triple in triples:
-        table.add_row(*triple)
-
-    return table
-
-
-def pairs_table(subject_predicate_pairs: List[tuple[str, str]], title: str) -> Table:
-
-    table = Table(title=title, show_lines=False)
-    table.add_column("Subject", justify="left")
-    table.add_column("Predicate", justify="left")
-    for pair in subject_predicate_pairs:
-        table.add_row(*pair)
-
-    return table

proscenium-0.0.8/proscenium/verbs/extract.py

@@ -1,65 +0,0 @@
-from typing import Optional
-import logging
-from rich.console import Console
-from string import Formatter
-
-import json
-from pydantic import BaseModel
-
-from proscenium.verbs.complete import complete_simple
-
-log = logging.getLogger(__name__)
-
-extraction_system_prompt = "You are an entity extractor"
-
-
-class PartialFormatter(Formatter):
-    def get_value(self, key, args, kwargs):
-        try:
-            return super().get_value(key, args, kwargs)
-        except KeyError:
-            return "{" + key + "}"
-
-
-partial_formatter = PartialFormatter()
-
-raw_extraction_template = """\
-Below is a description of a data class for storing information extracted from text:
-
-{extraction_description}
-
-Find the information in the following text, and provide them in the specified JSON response format.
-Only answer in JSON.:
-
-{text}
-"""
-
-
-def extract_to_pydantic_model(
-    extraction_model_id: str,
-    extraction_template: str,
-    clazz: type[BaseModel],
-    text: str,
-    console: Optional[Console] = None,
-) -> BaseModel:
-
-    extract_str = complete_simple(
-        extraction_model_id,
-        extraction_system_prompt,
-        extraction_template.format(text=text),
-        response_format={
-            "type": "json_object",
-            "schema": clazz.model_json_schema(),
-        },
-        console=console,
-    )
-
-    log.info("complete_to_pydantic_model: extract_str = <<<%s>>>", extract_str)
-
-    try:
-        extract_dict = json.loads(extract_str)
-        return clazz.model_construct(**extract_dict)
-    except Exception as e:
-        log.error("complete_to_pydantic_model: Exception: %s", e)
-
-    return None

proscenium-0.0.8/proscenium/verbs/read.py

@@ -1,53 +0,0 @@
-from typing import List
-
-import os
-import logging
-
-import httpx
-from pydantic.networks import HttpUrl
-from pathlib import Path
-
-from langchain_core.documents.base import Document
-from langchain_community.document_loaders import TextLoader
-from langchain_community.document_loaders.hugging_face_dataset import (
-    HuggingFaceDatasetLoader,
-)
-
-log = logging.getLogger(__name__)
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-logging.getLogger("langchain_text_splitters.base").setLevel(logging.ERROR)
-
-
-def load_file(filename: str) -> List[Document]:
-
-    loader = TextLoader(filename)
-    documents = loader.load()
-
-    return documents
-
-
-def load_hugging_face_dataset(
-    dataset_name: str, page_content_column: str = "text"
-) -> List[Document]:
-
-    loader = HuggingFaceDatasetLoader(
-        dataset_name, page_content_column=page_content_column
-    )
-    documents = loader.load()
-
-    return documents
-
-
-async def url_to_file(url: HttpUrl, data_file: Path, overwrite: bool = False):
-
-    if data_file.exists() and not overwrite:
-        return
-
-    async with httpx.AsyncClient() as client:
-
-        response = await client.get(url)
-        response.raise_for_status()
-
-        with open(data_file, "wb") as file:
-            file.write(response.content)

proscenium-0.0.8/proscenium/verbs/vector_database.py

@@ -1,139 +0,0 @@
-from typing import Dict, List
-
-import logging
-from pathlib import Path
-from langchain_core.documents.base import Document
-from urllib.parse import urlsplit
-from pymilvus import MilvusClient
-from pymilvus import DataType, FieldSchema, CollectionSchema
-from pymilvus import model
-
-# See https://milvus.io/docs/quickstart.md
-
-log = logging.getLogger(__name__)
-
-
-def embedding_function(
-    embedding_model_id: str,
-) -> model.dense.SentenceTransformerEmbeddingFunction:
-    embedding_fn = model.dense.SentenceTransformerEmbeddingFunction(
-        model_name=embedding_model_id, device="cpu"  # or 'cuda:0'
-    )
-    return embedding_fn
-
-
-def schema_chunks(
-    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
-) -> CollectionSchema:
-
-    field_id = FieldSchema(
-        name="id", dtype=DataType.INT64, is_primary=True, auto_id=True
-    )
-    field_text = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=50000)
-    field_vector = FieldSchema(
-        name="vector", dtype=DataType.FLOAT_VECTOR, dim=embedding_fn.dim
-    )
-
-    schema = CollectionSchema(
-        fields=[field_id, field_text, field_vector],
-        description="Chunks Schema",
-        enable_dynamic_field=True,
-    )
-
-    return schema
-
-
-def vector_db(
-    uri: str,
-) -> MilvusClient:
-
-    log.info("Connecting to vector db %s", uri)
-    uri_fields = urlsplit(uri)
-    client = None
-    if uri_fields[0] == "file":
-        file_path = Path(uri_fields[2][1:])
-        if file_path.exists():
-            log.info(
-                "Using existing %s file.",
-                uri_fields[2],
-            )
-        else:
-            log.info("Creating new vector db file %s", file_path)
-
-        client = MilvusClient(uri=str(file_path))
-
-    else:
-
-        log.info("Connecting to vector db at non-file uri %s", uri)
-        client = MilvusClient(uri=uri)
-
-    return client
-
-
-def create_collection(
-    client: MilvusClient,
-    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
-    collection_name: str,
-) -> None:
-
-    client.create_collection(
-        collection_name=collection_name,
-        schema=schema_chunks(embedding_fn),
-    )
-
-    index_params = client.prepare_index_params()
-
-    index_params.add_index(
-        field_name="vector",
-        index_type="IVF_FLAT",
-        metric_type="IP",
-        params={"nlist": 1024},
-    )
-
-    client.create_index(
-        collection_name=collection_name, index_params=index_params, sync=True
-    )
-    log.info("Created collection %s", collection_name)
-
-
-def add_chunks_to_vector_db(
-    client: MilvusClient,
-    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
-    chunks: List[Document],
-    collection_name: str,
-) -> Dict:
-
-    vectors = embedding_fn.encode_documents([chunk.page_content for chunk in chunks])
-
-    data = [
-        {"text": chunk.page_content, "vector": vector}
-        for chunk, vector in zip(chunks, vectors)
-    ]
-
-    insert_result = client.insert(collection_name, data)
-
-    return insert_result
-
-
-def closest_chunks(
-    client: MilvusClient,
-    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
-    query: str,
-    collection_name: str,
-    k: int = 4,
-) -> List[Dict]:
-
-    client.load_collection(collection_name)
-
-    result = client.search(
-        collection_name=collection_name,
-        data=embedding_fn.encode_queries([query]),
-        anns_field="vector",
-        search_params={"metric": "IP", "offset": 0},
-        output_fields=["text"],
-        limit=k,
-    )
-
-    hits = result[0]
-
-    return hits

proscenium-0.0.8/proscenium/verbs/write.py

@@ -1,14 +0,0 @@
-from typing import List
-
-import logging
-import csv
-
-log = logging.getLogger(__name__)
-
-
-def triples_to_csv(triples: List[tuple[str, str, str]], filename: str) -> None:
-
-    with open(filename, "wt") as f:
-        writer = csv.writer(f, delimiter=",", quotechar='"')
-        writer.writerow(["entity", "role", "case name"])  # header
-        writer.writerows(triples)