proscenium 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

proscenium/{scripts → patterns}/document_enricher.py

@@ -1,12 +1,13 @@
 from typing import List
 from typing import Callable
-from typing import Any
+from typing import Optional
 
 import time
+import logging
 from pydantic import BaseModel
 
-from rich import print
 from rich.panel import Panel
+from rich.console import Console
 from rich.progress import Progress
 
 from langchain_core.documents.base import Document
@@ -14,6 +15,8 @@ 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,
@@ -22,11 +25,12 @@ def extract_from_document_chunks(
     chunk_extraction_template: str,
     chunk_extract_clazz: type[BaseModel],
     delay: float,
-    verbose: bool = False,
+    console: Optional[Console] = None,
 ) -> List[BaseModel]:
 
-    print(doc_as_rich(doc))
-    print()
+    if console is not None:
+        console.print(doc_as_rich(doc))
+        console.print()
 
     extract_models = []
 
@@ -40,9 +44,9 @@ def extract_from_document_chunks(
             chunk.page_content,
         )
 
-        if verbose:
-            print("Extract model in chunk", i + 1, "of", len(chunks))
-            print(Panel(str(ce)))
+        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)
@@ -55,7 +59,7 @@ def enrich_documents(
     extract_from_doc_chunks: Callable[[Document], List[BaseModel]],
     doc_enrichments: Callable[[Document, list[BaseModel]], BaseModel],
     enrichments_jsonl_file: str,
-    verbose: bool = False,
+    console: Optional[Console] = None,
 ) -> None:
 
     docs = retrieve_documents()
@@ -70,11 +74,11 @@ def enrich_documents(
 
             for doc in docs:
 
-                chunk_extract_models = extract_from_doc_chunks(doc, verbose)
+                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)
 
-        print("Wrote document enrichments to", enrichments_jsonl_file)
+        log.info("Wrote document enrichments to %s", enrichments_jsonl_file)

proscenium/{scripts → patterns}/entity_resolver.py

@@ -1,6 +1,7 @@
 from typing import Optional
-from rich import print
+import logging
 
+from rich.console import Console
 from langchain_core.documents.base import Document
 from neo4j import Driver
 
@@ -13,6 +14,8 @@ 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:
 
@@ -21,27 +24,27 @@ class Resolver:
         cypher: str,
         field_name: str,
         collection_name: str,
-        embedding_model_id: str,
     ):
         self.cypher = cypher
         self.field_name = field_name
         self.collection_name = collection_name
-        self.embedding_model_id = embedding_model_id
 
 
 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, overwrite=True)
-    print("Vector db stored at", milvus_uri)
+    vector_db_client = vector_db(milvus_uri)
+    log.info("Vector db stored at %s", milvus_uri)
 
-    for resolver in resolvers:
+    embedding_fn = embedding_function(embedding_model_id)
+    log.info("Embedding model %s", embedding_model_id)
 
-        embedding_fn = embedding_function(resolver.embedding_model_id)
-        print("Embedding model", resolver.embedding_model_id)
+    for resolver in resolvers:
 
         values = []
         with driver.session() as session:
@@ -49,15 +52,16 @@ def load_entity_resolver(
             new_values = [Document(record[resolver.field_name]) for record in result]
             values.extend(new_values)
 
-        print("Loading entity resolver into vector db", resolver.collection_name)
-        create_collection(
-            vector_db_client, embedding_fn, resolver.collection_name, overwrite=True
-        )
+        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
         )
-        print(info["insert_count"], "chunks inserted")
-        print(collection_panel(vector_db_client, 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()
 
@@ -68,10 +72,12 @@ def find_matching_objects(
     resolver: Resolver,
 ) -> Optional[str]:
 
-    print("Loading collection", resolver.collection_name)
+    log.info("Loading collection", resolver.collection_name)
     vector_db_client.load_collection(resolver.collection_name)
 
-    print("Finding entity matches for", approximate, "using", resolver.collection_name)
+    log.info(
+        "Finding entity matches for", approximate, "using", resolver.collection_name
+    )
 
     hits = closest_chunks(
         vector_db_client,
@@ -82,8 +88,8 @@ def find_matching_objects(
     )
     # TODO apply distance threshold
     for match in [head["entity"]["text"] for head in hits[:1]]:
-        print("Closest match:", match)
+        log.info("Closest match:", match)
         return match
 
-    print("No match found")
+    log.info("No match found")
     return None

proscenium/patterns/graph_rag.py

@@ -0,0 +1,61 @@
+from typing import Callable
+from typing import Optional
+
+import logging
+
+from rich.console import Console
+
+from pydantic import BaseModel
+from uuid import uuid4, UUID
+from neo4j import Driver
+
+log = logging.getLogger(__name__)
+
+
+def query_to_prompts(
+    query: str,
+    query_extraction_model_id: str,
+    milvus_uri: str,
+    driver: Driver,
+    query_extract: Callable[
+        [str, str], BaseModel
+    ],  # (query_str, query_extraction_model_id) -> QueryExtractions
+    query_extract_to_graph: Callable[
+        [str, UUID, BaseModel], None
+    ],  # query, query_id, extract
+    query_extract_to_context: Callable[
+        [BaseModel, str, Driver, str, Optional[Console]], BaseModel
+    ],  # (QueryExtractions, query_str, Driver, milvus_uri) -> Context
+    context_to_prompts: Callable[
+        [BaseModel], tuple[str, str]
+    ],  # Context -> (system_prompt, user_prompt)
+    console: Optional[Console] = None,
+) -> Optional[tuple[str, str]]:
+
+    query_id = uuid4()
+
+    log.info("Extracting information from the question")
+
+    extract = query_extract(query, query_extraction_model_id)
+    if extract is None:
+        log.info("Unable to extract information from that question")
+        return None
+
+    log.info("Extract: %s", extract)
+
+    log.info("Storing the extracted information in the graph")
+    query_extract_to_graph(query, query_id, extract, driver)
+
+    log.info("Forming context from the extracted information")
+    context = query_extract_to_context(
+        extract, query, driver, milvus_uri, console=console
+    )
+    if context is None:
+        log.info("Unable to form context from the extracted information")
+        return None
+
+    log.info("Context: %s", context)
+
+    prompts = context_to_prompts(context)
+
+    return prompts

proscenium/{scripts → patterns}/knowledge_graph.py

@@ -1,14 +1,16 @@
 from typing import Callable
 from typing import Any
 
+import logging
 import json
 from pydantic import BaseModel
 
-from rich import print
 from rich.progress import Progress
 
 from neo4j import Driver
 
+log = logging.getLogger(__name__)
+
 
 def load_knowledge_graph(
     driver: Driver,
@@ -17,7 +19,7 @@ def load_knowledge_graph(
     doc_enrichments_to_graph: Callable[[Any, BaseModel], None],
 ) -> None:
 
-    print("Parsing enrichments from", enrichments_jsonl_file)
+    log.info("Parsing enrichments from %s", enrichments_jsonl_file)
 
     enrichmentss = []
     with open(enrichments_jsonl_file, "r") as f:

proscenium/{scripts → patterns}/rag.py

@@ -1,7 +1,5 @@
 from typing import List, Dict
-
-from rich import print
-from rich.panel import Panel
+import logging
 
 from pymilvus import MilvusClient
 from pymilvus import model
@@ -10,6 +8,7 @@ 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__)
 
 rag_system_prompt = "Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer."
 
@@ -44,20 +43,15 @@ def answer_question(
     vector_db_client: MilvusClient,
     embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
     collection_name: str,
-    verbose: bool = False,
 ) -> str:
 
-    print(Panel(query, title="User"))
-
     chunks = closest_chunks(vector_db_client, embedding_fn, query, collection_name)
-    if verbose:
-        print("Found", len(chunks), "closest chunks")
-        print(chunk_hits_table(chunks))
+    log.info("Found %s closest chunks", len(chunks))
+    log.info(chunk_hits_table(chunks))
 
     prompt = rag_prompt(chunks, query)
-    if verbose:
-        print("RAG prompt created. Calling inference at", model_id, "\n\n")
+    log.info("RAG prompt created. Calling inference at %s", model_id)
 
-    answer = complete_simple(model_id, rag_system_prompt, prompt, rich_output=verbose)
+    answer = complete_simple(model_id, rag_system_prompt, prompt)
 
     return answer

proscenium/{scripts → patterns}/tools.py

@@ -1,47 +1,17 @@
-from typing import List
+from typing import Optional
+import logging
 
-from rich import print
+from rich.console import Console
 from rich.panel import Panel
 from rich.text import Text
-from thespian.actors import Actor
-
-from gofannon.base import BaseTool
 
 from proscenium.verbs.complete import (
     complete_for_tool_applications,
     evaluate_tool_calls,
     complete_with_tool_results,
 )
-from proscenium.verbs.invoke import process_tools
-
-
-def tool_applier_actor_class(
-    tools: List[BaseTool],
-    system_message: str,
-    model_id: str,
-    temperature: float = 0.75,
-    rich_output: bool = False,
-):
-
-    tool_map, tool_desc_list = process_tools(tools)
 
-    class ToolApplier(Actor):
-
-        def receiveMessage(self, message, sender):
-
-            response = apply_tools(
-                model_id=model_id,
-                system_message=system_message,
-                message=message,
-                tool_desc_list=tool_desc_list,
-                tool_map=tool_map,
-                temperature=temperature,
-                rich_output=rich_output,
-            )
-
-            self.send(sender, response)
-
-    return ToolApplier
+log = logging.getLogger(__name__)
 
 
 def apply_tools(
@@ -51,7 +21,7 @@ def apply_tools(
     tool_desc_list: list,
     tool_map: dict,
     temperature: float = 0.75,
-    rich_output: bool = False,
+    console: Optional[Console] = None,
 ) -> str:
 
     messages = [
@@ -60,35 +30,33 @@ def apply_tools(
     ]
 
     response = complete_for_tool_applications(
-        model_id, messages, tool_desc_list, temperature, rich_output
+        model_id, messages, tool_desc_list, temperature, console
     )
 
     tool_call_message = response.choices[0].message
 
     if tool_call_message.tool_calls is None or len(tool_call_message.tool_calls) == 0:
 
-        if rich_output:
-            print(
+        if console is not None:
+            console.print(
                 Panel(
                     Text(str(tool_call_message.content)),
                     title="Tool Application Response",
                 )
             )
 
-        print("No tool applications detected")
+        log.info("No tool applications detected")
 
         return tool_call_message.content
 
     else:
 
-        if rich_output:
-            print(
+        if console is not None:
+            console.print(
                 Panel(Text(str(tool_call_message)), title="Tool Application Response")
             )
 
-        tool_evaluation_messages = evaluate_tool_calls(
-            tool_call_message, tool_map, rich_output
-        )
+        tool_evaluation_messages = evaluate_tool_calls(tool_call_message, tool_map)
 
         result = complete_with_tool_results(
             model_id,
@@ -97,7 +65,7 @@ def apply_tools(
             tool_evaluation_messages,
             tool_desc_list,
             temperature,
-            rich_output,
+            console,
         )
 
         return result

proscenium/verbs/__init__.py

@@ -0,0 +1,3 @@
+import logging
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())

proscenium/verbs/chunk.py

@@ -8,6 +8,8 @@ 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)
 

proscenium/verbs/complete.py

@@ -37,10 +37,12 @@ Valid model ids:
 - `ollama:granite3.1-dense:2b`
 """
 
+from typing import Optional
 from typing import Any
-
+import logging
 import json
-from rich import print
+
+from rich.console import Console
 from rich.console import Group
 from rich.panel import Panel
 from rich.table import Table
@@ -51,6 +53,8 @@ from aisuite.framework.message import ChatCompletionMessageToolCall
 
 from proscenium.verbs.display.tools import complete_with_tools_panel
 
+log = logging.getLogger(__name__)
+
 provider_configs = {
     # TODO expose this
     "ollama": {"timeout": 180},
@@ -63,14 +67,14 @@ def complete_simple(
     model_id: str, system_prompt: str, user_prompt: str, **kwargs
 ) -> str:
 
-    rich_output = kwargs.pop("rich_output", False)
+    console = kwargs.pop("console", None)
 
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_prompt},
     ]
 
-    if rich_output:
+    if console is not None:
 
         kwargs_text = "\n".join([str(k) + ": " + str(v) for k, v in kwargs.items()])
 
@@ -90,34 +94,30 @@ model_id: {model_id}
         call_panel = Panel(
             Group(params_text, messages_table), title="complete_simple call"
         )
-        print(call_panel)
+        console.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"))
+    if console is not None:
+        console.print(Panel(response, title="Response"))
 
     return response
 
 
-def evaluate_tool_call(
-    tool_map: dict, tool_call: ChatCompletionMessageToolCall, rich_output: bool = False
-) -> Any:
+def evaluate_tool_call(tool_map: dict, tool_call: ChatCompletionMessageToolCall) -> 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}")
+    log.info(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}")
+    log.info(f"   Response: {function_response}")
 
     return function_response
 
@@ -134,23 +134,19 @@ def tool_response_message(
     }
 
 
-def evaluate_tool_calls(
-    tool_call_message, tool_map: dict, rich_output: bool = False
-) -> list[dict]:
+def evaluate_tool_calls(tool_call_message, tool_map: dict) -> list[dict]:
 
     tool_call: ChatCompletionMessageToolCall
 
-    if rich_output:
-        print("Evaluating tool calls")
+    log.info("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)
+        function_response = evaluate_tool_call(tool_map, tool_call)
         new_messages.append(tool_response_message(tool_call, function_response))
 
-    if rich_output:
-        print("Tool calls evaluated")
+    log.info("Tool calls evaluated")
 
     return new_messages
 
@@ -160,10 +156,10 @@ def complete_for_tool_applications(
     messages: list,
     tool_desc_list: list,
     temperature: float,
-    rich_output: bool = False,
+    console: Optional[Console] = None,
 ):
 
-    if rich_output:
+    if console is not None:
         panel = complete_with_tools_panel(
             "complete for tool applications",
             model_id,
@@ -171,7 +167,7 @@ def complete_for_tool_applications(
             messages,
             temperature,
         )
-        print(panel)
+        console.print(panel)
 
     response = client.chat.completions.create(
         model=model_id,
@@ -190,13 +186,13 @@ def complete_with_tool_results(
     tool_evaluation_messages: list[dict],
     tool_desc_list: list,
     temperature: float,
-    rich_output: bool = False,
+    console: Optional[Console] = None,
 ):
 
     messages.append(tool_call_message)
     messages.extend(tool_evaluation_messages)
 
-    if rich_output:
+    if console is not None:
         panel = complete_with_tools_panel(
             "complete call with tool results",
             model_id,
@@ -204,7 +200,7 @@ def complete_with_tool_results(
             messages,
             temperature,
         )
-        print(panel)
+        console.print(panel)
 
     response = client.chat.completions.create(
         model=model_id,

proscenium/verbs/display/__init__.py

@@ -4,6 +4,6 @@ from rich.text import Text
 def header() -> Text:
     text = Text()
     text.append("Proscenium 🎭\n", style="bold")
-    text.append("The AI Alliance\n", style="bold")
+    text.append("https://the-ai-alliance.github.io/proscenium/\n")
     # TODO version, timestamp, ...
     return text

proscenium/verbs/display.py

@@ -1,5 +1,8 @@
+import logging
 from rich.text import Text
 
+log = logging.getLogger(__name__)
+
 
 def header() -> Text:
     text = Text(

proscenium/verbs/extract.py

@@ -1,4 +1,6 @@
+from typing import Optional
 import logging
+from rich.console import Console
 from string import Formatter
 
 import json
@@ -6,6 +8,8 @@ from pydantic import BaseModel
 
 from proscenium.verbs.complete import complete_simple
 
+log = logging.getLogger(__name__)
+
 extraction_system_prompt = "You are an entity extractor"
 
 
@@ -36,7 +40,7 @@ def extract_to_pydantic_model(
     extraction_template: str,
     clazz: type[BaseModel],
     text: str,
-    verbose: bool = False,
+    console: Optional[Console] = None,
 ) -> BaseModel:
 
     extract_str = complete_simple(
@@ -47,15 +51,15 @@ def extract_to_pydantic_model(
             "type": "json_object",
             "schema": clazz.model_json_schema(),
         },
-        rich_output=verbose,
+        console=console,
     )
 
-    logging.info("complete_to_pydantic_model: extract_str = <<<%s>>>", extract_str)
+    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:
-        logging.error("complete_to_pydantic_model: Exception: %s", e)
+        log.error("complete_to_pydantic_model: Exception: %s", e)
 
     return None

proscenium/verbs/invoke.py

@@ -1,5 +1,8 @@
+import logging
 from gofannon.base import BaseTool
 
+log = logging.getLogger(__name__)
+
 
 def process_tools(tools: list[BaseTool]) -> tuple[dict, list]:
     applied_tools = [F() for F in tools]