proscenium 0.0.1__py3-none-any.whl

proscenium/scripts/chunk_space.py

@@ -0,0 +1,33 @@
+from rich import print
+
+from pymilvus import MilvusClient
+from pymilvus import model
+
+from proscenium.verbs.read import load_file
+from proscenium.verbs.chunk import documents_to_chunks_by_characters
+from proscenium.verbs.vector_database import create_collection
+from proscenium.verbs.vector_database import add_chunks_to_vector_db
+from proscenium.verbs.display.milvus import collection_panel
+
+
+def build_vector_db(
+    data_files: list[str],
+    vector_db_client: MilvusClient,
+    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
+    collection_name: str,
+):
+
+    create_collection(vector_db_client, embedding_fn, collection_name, overwrite=True)
+
+    for data_file in data_files:
+
+        documents = load_file(data_file)
+        chunks = documents_to_chunks_by_characters(documents)
+        print("Data file", data_file, "has", len(chunks), "chunks")
+
+        info = add_chunks_to_vector_db(
+            vector_db_client, embedding_fn, chunks, collection_name
+        )
+        print(info["insert_count"], "chunks inserted")
+
+    print(collection_panel(vector_db_client, collection_name))

proscenium/scripts/document_enricher.py

@@ -0,0 +1,80 @@
+from typing import List
+from typing import Callable
+from typing import Any
+
+import time
+from pydantic import BaseModel
+
+from rich import print
+from rich.panel import Panel
+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
+
+
+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,
+    verbose: bool = False,
+) -> List[BaseModel]:
+
+    print(doc_as_rich(doc))
+    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,
+        )
+
+        if verbose:
+            print("Extract model in chunk", i + 1, "of", len(chunks))
+            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,
+    verbose: bool = False,
+) -> 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, verbose)
+                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)

proscenium/scripts/entity_resolver.py

@@ -0,0 +1,89 @@
+from typing import Optional
+from rich import print
+
+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
+
+
+class Resolver:
+
+    def __init__(
+        self,
+        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],
+    milvus_uri: str,
+) -> None:
+
+    vector_db_client = vector_db(milvus_uri, overwrite=True)
+    print("Vector db stored at", milvus_uri)
+
+    for resolver in resolvers:
+
+        embedding_fn = embedding_function(resolver.embedding_model_id)
+        print("Embedding model", resolver.embedding_model_id)
+
+        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)
+
+        print("Loading entity resolver into vector db", resolver.collection_name)
+        create_collection(
+            vector_db_client, embedding_fn, resolver.collection_name, overwrite=True
+        )
+        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))
+
+    vector_db_client.close()
+
+
+def find_matching_objects(
+    vector_db_client: MilvusClient,
+    approximate: str,
+    resolver: Resolver,
+) -> Optional[str]:
+
+    print("Loading collection", resolver.collection_name)
+    vector_db_client.load_collection(resolver.collection_name)
+
+    print("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]]:
+        print("Closest match:", match)
+        return match
+
+    print("No match found")
+    return None

proscenium/scripts/graph_rag.py

@@ -0,0 +1,43 @@
+from typing import Callable
+
+from pydantic import BaseModel
+
+from rich import print
+from rich.panel import Panel
+
+from neo4j import Driver
+
+from proscenium.verbs.complete import complete_simple
+
+
+def query_to_prompts(
+    question: str,
+    query_extraction_model_id: str,
+    milvus_uri: str,
+    driver: Driver,
+    query_extract: Callable[
+        [str, str, bool], BaseModel
+    ],  # (query_str, query_extraction_model_id) -> QueryExtractions
+    extract_to_context: Callable[
+        [BaseModel, str, Driver, str, bool], BaseModel
+    ],  # (QueryExtractions, query_str, Driver, milvus_uri) -> Context
+    context_to_prompts: Callable[
+        [BaseModel, bool], tuple[str, str]
+    ],  # Context -> (system_prompt, user_prompt)
+    verbose: bool = False,
+) -> str:
+
+    print("Extracting information from the question")
+    extract = query_extract(question, query_extraction_model_id, verbose)
+    if extract is None:
+        print("Unable to extract information from that question")
+        return None
+    print("Extract:", extract)
+
+    print("Forming context from the extracted information")
+    context = extract_to_context(extract, question, driver, milvus_uri, verbose)
+    print("Context:", context)
+
+    prompts = context_to_prompts(context, verbose)
+
+    return prompts

proscenium/scripts/knowledge_graph.py

@@ -0,0 +1,39 @@
+from typing import Callable
+from typing import Any
+
+import json
+from pydantic import BaseModel
+
+from rich import print
+from rich.progress import Progress
+
+from neo4j import Driver
+
+
+def load_knowledge_graph(
+    driver: Driver,
+    enrichments_jsonl_file: str,
+    enrichments_clazz: type[BaseModel],
+    doc_enrichments_to_graph: Callable[[Any, BaseModel], None],
+) -> None:
+
+    print("Parsing enrichments from", 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/scripts/rag.py

@@ -0,0 +1,63 @@
+from typing import List, Dict
+
+from rich import print
+from rich.panel import Panel
+
+from pymilvus import MilvusClient
+from pymilvus import model
+
+from proscenium.verbs.complete import complete_simple
+from proscenium.verbs.display.milvus import chunk_hits_table
+from proscenium.verbs.vector_database import closest_chunks
+
+
+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."
+
+rag_prompt_template = """
+The document chunks that are most similar to the query are:
+
+{context}
+
+Question:
+
+{query}
+
+Answer:
+"""
+
+
+def rag_prompt(chunks: List[Dict], query: str) -> str:
+
+    context = "\n\n".join(
+        [
+            f"CHUNK {chunk['id']}. {chunk['entity']['text']}"
+            for i, chunk in enumerate(chunks)
+        ]
+    )
+
+    return rag_prompt_template.format(context=context, query=query)
+
+
+def answer_question(
+    query: str,
+    model_id: str,
+    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))
+
+    prompt = rag_prompt(chunks, query)
+    if verbose:
+        print("RAG prompt created. Calling inference at", model_id, "\n\n")
+
+    answer = complete_simple(model_id, rag_system_prompt, prompt, rich_output=verbose)
+
+    return answer

proscenium/scripts/tools.py

@@ -0,0 +1,103 @@
+from typing import List
+
+from rich import print
+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
+
+
+def apply_tools(
+    model_id: str,
+    system_message: str,
+    message: str,
+    tool_desc_list: list,
+    tool_map: dict,
+    temperature: float = 0.75,
+    rich_output: bool = False,
+) -> str:
+
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": message},
+    ]
+
+    response = complete_for_tool_applications(
+        model_id, messages, tool_desc_list, temperature, rich_output
+    )
+
+    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(
+                Panel(
+                    Text(str(tool_call_message.content)),
+                    title="Tool Application Response",
+                )
+            )
+
+        print("No tool applications detected")
+
+        return tool_call_message.content
+
+    else:
+
+        if rich_output:
+            print(
+                Panel(Text(str(tool_call_message)), title="Tool Application Response")
+            )
+
+        tool_evaluation_messages = evaluate_tool_calls(
+            tool_call_message, tool_map, rich_output
+        )
+
+        result = complete_with_tool_results(
+            model_id,
+            messages,
+            tool_call_message,
+            tool_evaluation_messages,
+            tool_desc_list,
+            temperature,
+            rich_output,
+        )
+
+        return result

proscenium/verbs/chunk.py

@@ -0,0 +1,40 @@
+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
+
+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