proscenium 0.0.1__py3-none-any.whl

proscenium/verbs/complete.py

@@ -0,0 +1,214 @@
+"""
+This module uses the [`aisuite`](https://github.com/andrewyng/aisuite) library
+to interact with various LLM inference providers.
+
+It provides functions to complete a simple chat prompt, evaluate a tool call,
+and apply a list of tool calls to a chat prompt.
+
+Providers tested with Proscenium include:
+
+# AWS
+
+Environment: `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`
+
+Valid model ids:
+- `aws:meta.llama3-1-8b-instruct-v1:0`
+
+# Anthropic
+
+Environment: `ANTHROPIC_API_KEY`
+
+Valid model ids:
+- `anthropic:claude-3-5-sonnet-20240620`
+
+# OpenAI
+
+Environment: `OPENAI_API_KEY`
+
+Valid model ids:
+- `openai:gpt-4o`
+
+# Ollama
+
+Command line, eg `ollama run llama3.2 --keepalive 2h`
+
+Valid model ids:
+- `ollama:llama3.2`
+- `ollama:granite3.1-dense:2b`
+"""
+
+from typing import Any
+
+import json
+from rich import print
+from rich.console import Group
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from aisuite import Client
+from aisuite.framework.message import ChatCompletionMessageToolCall
+
+from proscenium.verbs.display.tools import complete_with_tools_panel
+
+provider_configs = {
+    # TODO expose this
+    "ollama": {"timeout": 180},
+}
+
+client = Client(provider_configs=provider_configs)
+
+
+def complete_simple(
+    model_id: str, system_prompt: str, user_prompt: str, **kwargs
+) -> str:
+
+    rich_output = kwargs.pop("rich_output", False)
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+
+    if rich_output:
+
+        kwargs_text = "\n".join([str(k) + ": " + str(v) for k, v in kwargs.items()])
+
+        params_text = Text(
+            f"""
+model_id: {model_id}
+{kwargs_text}
+    """
+        )
+
+        messages_table = Table(title="Messages", show_lines=True)
+        messages_table.add_column("Role", justify="left")
+        messages_table.add_column("Content", justify="left")  # style="green"
+        for message in messages:
+            messages_table.add_row(message["role"], message["content"])
+
+        call_panel = Panel(
+            Group(params_text, messages_table), title="complete_simple call"
+        )
+        print(call_panel)
+
+    response = client.chat.completions.create(
+        model=model_id, messages=messages, **kwargs
+    )
+    response = response.choices[0].message.content
+
+    if rich_output:
+        print(Panel(response, title="Response"))
+
+    return response
+
+
+def evaluate_tool_call(
+    tool_map: dict, tool_call: ChatCompletionMessageToolCall, rich_output: bool = False
+) -> Any:
+
+    function_name = tool_call.function.name
+    # TODO validate the arguments?
+    function_args = json.loads(tool_call.function.arguments)
+
+    if rich_output:
+        print(f"Evaluating tool call: {function_name} with args {function_args}")
+
+    function_response = tool_map[function_name](**function_args)
+
+    if rich_output:
+        print(f"   Response: {function_response}")
+
+    return function_response
+
+
+def tool_response_message(
+    tool_call: ChatCompletionMessageToolCall, tool_result: Any
+) -> dict:
+
+    return {
+        "role": "tool",
+        "tool_call_id": tool_call.id,
+        "name": tool_call.function.name,
+        "content": json.dumps(tool_result),
+    }
+
+
+def evaluate_tool_calls(
+    tool_call_message, tool_map: dict, rich_output: bool = False
+) -> list[dict]:
+
+    tool_call: ChatCompletionMessageToolCall
+
+    if rich_output:
+        print("Evaluating tool calls")
+
+    new_messages: list[dict] = []
+
+    for tool_call in tool_call_message.tool_calls:
+        function_response = evaluate_tool_call(tool_map, tool_call, rich_output)
+        new_messages.append(tool_response_message(tool_call, function_response))
+
+    if rich_output:
+        print("Tool calls evaluated")
+
+    return new_messages
+
+
+def complete_for_tool_applications(
+    model_id: str,
+    messages: list,
+    tool_desc_list: list,
+    temperature: float,
+    rich_output: bool = False,
+):
+
+    if rich_output:
+        panel = complete_with_tools_panel(
+            "complete for tool applications",
+            model_id,
+            tool_desc_list,
+            messages,
+            temperature,
+        )
+        print(panel)
+
+    response = client.chat.completions.create(
+        model=model_id,
+        messages=messages,
+        temperature=temperature,
+        tools=tool_desc_list,  # tool_choice="auto",
+    )
+
+    return response
+
+
+def complete_with_tool_results(
+    model_id: str,
+    messages: list,
+    tool_call_message: dict,
+    tool_evaluation_messages: list[dict],
+    tool_desc_list: list,
+    temperature: float,
+    rich_output: bool = False,
+):
+
+    messages.append(tool_call_message)
+    messages.extend(tool_evaluation_messages)
+
+    if rich_output:
+        panel = complete_with_tools_panel(
+            "complete call with tool results",
+            model_id,
+            tool_desc_list,
+            messages,
+            temperature,
+        )
+        print(panel)
+
+    response = client.chat.completions.create(
+        model=model_id,
+        messages=messages,
+    )
+
+    return response.choices[0].message.content

proscenium/verbs/display/__init__.py

@@ -0,0 +1,9 @@
+from rich.text import Text
+
+
+def header() -> Text:
+    text = Text()
+    text.append("Proscenium 🎭\n", style="bold")
+    text.append("The AI Alliance\n", style="bold")
+    # TODO version, timestamp, ...
+    return text

proscenium/verbs/display/chat.py

@@ -0,0 +1,35 @@
+from rich.table import Table
+
+
+def messages_table(messages: list) -> Table:
+
+    table = Table(title="Messages in Chat Context", show_lines=True)
+    table.add_column("Role", justify="left")
+    table.add_column("Content", justify="left")
+    for message in messages:
+        if type(message) is dict:
+            role = message["role"]
+            content = ""
+            if role == "tool":
+                content = f"""tool call id = {message['tool_call_id']}
+fn name = {message['name']}
+result = {message['content']}"""
+            elif role == "assistant":
+                content = f"""{str(message)}"""
+            else:
+                content = message["content"]
+            table.add_row(role, content)
+        else:
+            role = message.role
+            content = ""
+            if role == "tool":
+                content = f"""tool call id = {message.tool_call_id}
+fn name = {message.name}
+result = {message['content']}"""
+            elif role == "assistant":
+                content = f"""{str(message)}"""
+            else:
+                content = message.content
+            table.add_row(role, content)
+
+    return table

proscenium/verbs/display/milvus.py

@@ -0,0 +1,68 @@
+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/verbs/display/neo4j.py

@@ -0,0 +1,25 @@
+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/verbs/display/tools.py

@@ -0,0 +1,64 @@
+from rich.console import Group
+from rich.table import Table
+from rich.text import Text
+from rich.panel import Panel
+
+from .chat import messages_table
+
+
+def parameters_table(parameters: list[dict]) -> Table:
+
+    table = Table(title="Parameters", show_lines=False, box=None)
+    table.add_column("name", justify="right")
+    table.add_column("type", justify="left")
+    table.add_column("description", justify="left")
+
+    for name, props in parameters["properties"].items():
+        table.add_row(name, props["type"], props["description"])
+
+    # TODO denote required params
+
+    return table
+
+
+def function_description_panel(fd: dict) -> Panel:
+
+    fn = fd["function"]
+
+    text = Text(f"{fd['type']} {fn['name']}: {fn['description']}\n")
+
+    pt = parameters_table(fn["parameters"])
+
+    panel = Panel(Group(text, pt))
+
+    return panel
+
+
+def function_descriptions_panel(function_descriptions: list[dict]) -> Panel:
+
+    sub_panels = [function_description_panel(fd) for fd in function_descriptions]
+
+    panel = Panel(Group(*sub_panels), title="Function Descriptions")
+
+    return panel
+
+
+def complete_with_tools_panel(
+    title: str, model_id: str, tool_desc_list: list, messages: list, temperature: float
+) -> Panel:
+
+    text = Text(
+        f"""
+model_id: {model_id}
+temperature: {temperature}
+"""
+    )
+
+    panel = Panel(
+        Group(
+            text, function_descriptions_panel(tool_desc_list), messages_table(messages)
+        ),
+        title=title,
+    )
+
+    return panel

proscenium/verbs/display.py

@@ -0,0 +1,10 @@
+from rich.text import Text
+
+
+def header() -> Text:
+    text = Text(
+        """[bold]Proscenium[/bold] :performing_arts:
+[bold]The AI Alliance[/bold]"""
+    )
+    # TODO version, timestamp, ...
+    return text

proscenium/verbs/extract.py

@@ -0,0 +1,61 @@
+import logging
+from string import Formatter
+
+import json
+from pydantic import BaseModel
+
+from proscenium.verbs.complete import complete_simple
+
+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,
+    verbose: bool = False,
+) -> 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(),
+        },
+        rich_output=verbose,
+    )
+
+    logging.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:
+        logging.error("complete_to_pydantic_model: Exception: %s", e)
+
+    return None

proscenium/verbs/invoke.py

@@ -0,0 +1,8 @@
+from gofannon.base import BaseTool
+
+
+def process_tools(tools: list[BaseTool]) -> tuple[dict, list]:
+    applied_tools = [F() for F in tools]
+    tool_map = {f.name: f.fn for f in applied_tools}
+    tool_desc_list = [f.definition for f in applied_tools]
+    return tool_map, tool_desc_list

proscenium/verbs/know.py

@@ -0,0 +1,9 @@
+from neo4j import GraphDatabase
+from neo4j import Driver
+
+
+def knowledge_graph_client(uri: str, username: str, password: str) -> Driver:
+
+    driver = GraphDatabase.driver(uri, auth=(username, password))
+
+    return driver

proscenium/verbs/read.py

@@ -0,0 +1,55 @@
+from typing import List
+
+import os
+import logging
+
+from langchain_core.documents.base import Document
+
+from langchain_community.document_loaders import TextLoader
+from langchain_community.document_loaders.hugging_face_dataset import (
+    HuggingFaceDatasetLoader,
+)
+
+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
+
+
+import httpx
+from pydantic.networks import HttpUrl
+from pathlib import Path
+
+
+async def url_to_file(url: HttpUrl, data_file: Path, overwrite: bool = False):
+
+    if data_file.exists() and not overwrite:
+        # print(f"File {data_file} exists. Use overwrite=True to replace.")
+        return
+
+    async with httpx.AsyncClient() as client:
+
+        # print(f"Downloading {url} to {data_file}...")
+        response = await client.get(url)
+        response.raise_for_status()
+
+        with open(data_file, "wb") as file:
+            file.write(response.content)

proscenium/verbs/remember.py

@@ -0,0 +1,8 @@
+def format_chat_history(chat_history) -> str:
+    delimiter = "-" * 80 + "\n"
+    return delimiter.join(
+        [
+            f"{msg['sender']} to {msg['receiver']}:\n\n{msg['content']}\n\n"
+            for msg in chat_history
+        ]
+    )