Tools update (#21)

neumai 0.0.33
- Small updates to website connector
- Addition of vector as parameters on search result (only added for Weaviate, but other should be trivial)
- Pipeline description field to be used for routed based search

neumai-tool 0.0.17
- PipelineCollections
- Dataset evaluation
- Interop helpers
- Update of semantic helpers

Author: ddematheu

neumai-tools/README.md

@@ -1,3 +1,13 @@
 # Neum AI Tools
 
-Collection of tools made to help implement RAG pipelines. Can be used directly or through Neum AI.
+Collection of tools made to help implement RAG pipelines. Can be used directly or through Neum AI.
+
+| :exclamation:  This utilities are currently in experimental phase.   |
+|----------------------------------------------------------------------|
+
+## Tools
+
+- [Semantic Helpers](./neumai_tools/SemanticHelpers/): LLM based tools to help augment RAG pipelines. Includes generating semantic strategies to generate chunking code, select fields to capture as metadata and as content when leveraging RAG on top of structured data.
+- [Interop Helpers](./neumai_tools/InteropHelpers/): Utilities to connect frameworks like Langchain and Llama Index with Neum AI. If there are data connectors that you want to leverage, you can use the utilities to translate interfaces.
+- [Pipeline Collection](./neumai_tools/PipelineCollection/): Treat a collection of pipelines as a single entity to perform search and other actions. Allows you to silo data into pipelines with their unique set of transformations, embed model and sink and then bring them back together as a single entity to retrieve data.
+- [Dataset Evaluation](./neumai_tools/DatasetEvaluation/): Create datasets of queries and expected outputs that you can run against a given pipeline or pipeline collection. 

neumai-tools/neumai_tools/DatasetEvaluation/Dataset.py

@@ -0,0 +1,59 @@
+from pydantic import BaseModel, Field
+from neumai.Pipelines.Pipeline import Pipeline
+from neumai_tools.PipelineCollection.PipelineCollection import PipelineCollection
+from .Evaluation import Evaluation, CosineEvaluation
+from .DatasetUtils import DatasetEntry, DatasetResult, DatasetResults
+from uuid import uuid4
+from typing import List, Type
+import numpy as np
+
+class Dataset(BaseModel):
+
+    name : str = Field(... , description="")
+    dataset_entries : List[DatasetEntry] = Field(... , description="")
+    evaluation_type: Type[Evaluation] = Field(...,description="Evaluation type for results.")
+
+    def run_with_pipeline(self, pipeline: Pipeline) -> DatasetResults:
+
+        dataset_results = DatasetResults()
+
+        for dataset_entry in self.dataset_entries:
+            # Generate results
+            result_output = pipeline.search(query=dataset_entry.query, number_of_results=1)[0]
+
+            dataset_result = self.evaluation_type(pipeline=pipeline, dataset_entry=dataset_entry, result_output=result_output).evaluate()
+
+            dataset_results.dataset_results.append(dataset_result)
+        
+        return dataset_results
+
+    
+    def run_with_pipeline_collection_unified(self, pipeline_collection: PipelineCollection):
+        dataset_results = DatasetResults()
+
+        for dataset_entry in self.dataset_entries:
+            # Generate results
+            result_output = pipeline_collection.search_unified(query=dataset_entry.query, number_of_results=1)[0]
+
+            dataset_result = self.evaluation_type(pipeline=pipeline_collection.pipelines[0], dataset_entry=dataset_entry, result_output=result_output).evaluate()
+
+            dataset_results.dataset_results.append(dataset_result)
+        
+        return dataset_results
+    
+    def run_with_pipeline_collection_separate(self, pipeline_collection: PipelineCollection):
+        dataset_results_separate: dict[str,DatasetResults] = {}
+        for pipeline in pipeline_collection.pipelines:
+            dataset_results = DatasetResults()
+
+            for dataset_entry in self.dataset_entries:
+                # Generate results
+                result_output = pipeline.search(query=dataset_entry.query, number_of_results=1)[0]
+                
+                dataset_result = self.evaluation_type(pipeline=pipeline, dataset_entry=dataset_entry, result_output=result_output).evaluate()
+
+                dataset_results.dataset_results.append(dataset_result)
+
+            dataset_results_separate[pipeline.id] = dataset_results
+            
+        return dataset_results_separate

neumai-tools/neumai_tools/DatasetEvaluation/DatasetUtils.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel, Field
+from uuid import uuid4
+from neumai.Shared.NeumSearch import NeumSearchResult
+from typing import List, Optional
+
+class DatasetEntry(BaseModel):
+
+    id : str =  Field(uuid4() , description="")
+    query : str = Field(... , description="")
+    expected_output : str = Field(... , description="")
+
+class DatasetResult(BaseModel):
+
+    dataset_entry : DatasetEntry = Field(... , description="")
+    raw_result : NeumSearchResult = Field(... , description="")
+    score : Optional[float] = Field(None, description="")
+    evaluation: Optional[dict] =  Field(None, description="")
+
+class DatasetResults(BaseModel):
+    dataset_results_id : str = Field(uuid4(), description="")
+    dataset_results : List[DatasetResult] = Field([], description="")

neumai-tools/neumai_tools/DatasetEvaluation/Evaluation.py

@@ -0,0 +1,125 @@
+from pydantic import BaseModel, Field
+from abc import ABC, abstractmethod
+from neumai.Pipelines.Pipeline import Pipeline
+from .DatasetUtils import DatasetResult, DatasetEntry
+from neumai.Shared.NeumSearch import NeumSearchResult
+import numpy as np
+import json
+import re
+
+class Evaluation(BaseModel, ABC):
+    pipeline: Pipeline = Field(...,description="")
+    dataset_entry: DatasetEntry = Field(...,description="Dataset entry for evaluation")
+    result_output: NeumSearchResult = Field(..., description="Actual result for retrieval")
+
+    @abstractmethod
+    def evaluate(self) -> DatasetResult:
+        """Method that runs evaluation and generates results"""
+
+class CosineEvaluation(Evaluation):
+
+    def evaluate(self) -> DatasetResult:
+        # Calculate score -> cosine similarity between expected output and result
+
+        # Calculate vectors for both outputs
+        expected_output_vector = self.pipeline.embed.embed_query(query=self.dataset_entry.expected_output)
+        result_vector = self.pipeline.embed.embed_query(query=self.result_output.metadata['text'])
+
+        # Normalize the vectors to unit vectors
+        expected_output_vector_norm = expected_output_vector / np.linalg.norm(expected_output_vector)
+        result_vector_norm = result_vector / np.linalg.norm(result_vector)
+
+        # Calculate the dot product and return
+        similarity = np.dot(expected_output_vector_norm, result_vector_norm)
+
+        dataset_result = DatasetResult(
+            dataset_entry=self.dataset_entry, 
+            raw_result=self.result_output,
+            score=similarity
+        )
+
+        return dataset_result
+
+class LLMEvaluation(Evaluation):
+
+    def evaluate(self) -> DatasetResult:
+        system_prompt = """
+        Evaluate the following questions and pieces of context. 
+
+        Question: Was ketchup originally a type of medicine?
+        Expected output: Yes, in the 1830's ketchup was sold as a medicine.
+        Actual output: Ketchup was sold in the 1830s as medicine. In 1834, it was sold as a cure for an upset stomach by an Ohio physician named John Cook. It wasn't popularized as a condiment until the late 19th century!
+        Evaluation: {
+            "relevant_context": true,
+            "all_context_present": true
+        }
+
+        Question: Where did the shortest war in history happen?
+        Expected output: The war took place in Zanzibar.
+        Actual output: The shortest war in history lasted 38 minutes! It was between Britain and Zanzibar and is known as the Anglo-Zanzibar War. This war occurred on August 27, 1896. It was over the ascension of the next Sultan in Zanzibar and resulted in a British victory.
+        Evaluation: {
+            "relevant_context": true,
+            "all_context_present": false
+        }
+
+        Question: What animals did Roman's have as pets?
+        Expected output: Ferrets, dogs and monkeys were the most popular pets in the Roman Empire.
+        Actual output: Roman's didn't have cats as pets.
+        Evaluation: {
+            "relevant_context": false,
+            "all_context_present": false
+        }
+        """
+
+
+        user_prompt = """Question: {question}
+        Expected output: {expected_output}
+        Actual output: {actual_output}
+        Evaluation:
+        """
+
+        generated_user_prompt = user_prompt.format(
+            question=self.dataset_entry.query,
+            expected_output = self.dataset_entry.expected_output,
+            actual_output = self.result_output.metadata['text']
+        )
+        
+        from openai import OpenAI
+        client = OpenAI()
+
+        response = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {
+                    "role": "system",
+                    "content": system_prompt
+                },
+                {
+                    "role":"user",
+                    "content":generated_user_prompt
+                }
+            ],
+            temperature=1,
+        )
+        
+        match = re.search(r'\{.*?\}', response.choices[0].message.content, re.DOTALL)
+
+        if match:
+            json_object = match.group()
+            # Now, json_object contains the JSON string.
+            # To convert it to a Python dictionary, use json.loads if it's a valid JSON
+            try:
+                evaluation = json.loads(json_object)
+                # data is now a Python dictionary representing your JSON object
+            except json.JSONDecodeError:
+                print("The extracted string is not a valid JSON.")
+                evaluation = None
+        else:
+            print("No JSON object found in the string.")
+            evaluation = None
+
+        return DatasetResult(
+            dataset_entry=self.dataset_entry,
+            raw_result=self.result_output,
+            evaluation=evaluation
+        )

neumai-tools/neumai_tools/DatasetEvaluation/__init__.py

@@ -0,0 +1,18 @@
+from langchain.docstore.document import Document as LangChainDocument
+from llama_index import Document as LlamaIndexDocument
+from neumai.Shared.NeumDocument import NeumDocument
+from uuid import uuid4
+
+def document_transformer_langchain(document:LangChainDocument, id:str = uuid4()):
+    return NeumDocument(
+        id=id,
+        content=document.page_content,
+        metadata=document.metadata
+    )
+
+def document_transformer_llamaIndex(document:LlamaIndexDocument):
+    return NeumDocument(
+        id=document.doc_id,
+        content=document.text,
+        metadata=document.metadata
+    )

neumai-tools/neumai_tools/InteropHelpers/DocumentTransformer.py

@@ -0,0 +1,4 @@
+from .DocumentTransformer import (
+    document_transformer_langchain,
+    document_transformer_llamaIndex
+)

neumai-tools/neumai_tools/InteropHelpers/__init__.py

@@ -0,0 +1,56 @@
+from typing import List
+from pydantic import BaseModel, Field
+from neumai.Pipelines.Pipeline import Pipeline
+from neumai.Shared.NeumSearch import NeumSearchResult
+
+class PipelineCollection(BaseModel):
+    """
+    Pipeline Collection
+
+    Collections of Pipelines that can be managed as a unit. Includes running and searching them as a unit.
+
+    Attributes
+    ----------
+
+    pipelines : List[Pipeline]
+        List of Neum AI Pipelines
+    
+    """
+    
+    pipelines: List[Pipeline] = Field(..., description = "Pipelines that are to be par of the collection.")
+
+    def run(self):
+        for pipeline in self.pipelines:
+            pipeline.run()
+    
+    def search_unified(self, query:str, number_of_results:int) -> List[NeumSearchResult]:
+        """Search pipelines and unify results using scores as re-ranking criteria"""
+        
+        # Simple query of results
+        search_results = []
+        for pipeline in self.pipelines:
+            results = pipeline.search(query=query, number_of_results=number_of_results)
+            search_results.append(results)
+
+        # Re-rank results by score. Should add a callback option for re-rank
+        return sorted(search_results, key=lambda x: x.score, reverse=True)[:number_of_results]
+
+    def search_separate(self, query:str, number_of_results:int)-> List:
+        """Search pipelines and provide raw results for each pipeline"""
+        
+        # Simple query of results
+        search_results = []
+        for pipeline in self.pipelines:
+            results = pipeline.search(query=query, number_of_results=number_of_results)
+            search_results.append({
+                "pipeline_id":pipeline.id,
+                "results" : results
+            })
+
+        # Re-rank results by score. Should add a callback option for re-rank
+        return (search_results)
+    
+    def search_routed(self, query:str, number_of_results:int)-> List:
+        """Routed search based on the contents available in a pipeline"""
+        # Need to add descriptions to the pipeline and generate a basic index on top of them
+        raise NotImplementedError("In the works. Contact founders@tryneum.com for information")

neumai-tools/neumai_tools/PipelineCollection/PipelineCollection.py

@@ -1,5 +1,5 @@
 import openai
-from langchain.docstore.document import Document
+from neumai.Shared.NeumDocument import NeumDocument
 from typing import (
     List,
 )
@@ -32,7 +32,7 @@ def semantic_chunking_strategy_code(text:str, chunking_strategy:str) -> str:
     messages = [
         {"role": "system", "content": ('You are helpful developer that writes python code.' + 
                                         'Output the code in this format: ```python def split_text_into_chunks(text): <Insert Code>```' +
-                                        'The function `split_text_into_chunks` should output an array of chunks.'
+                                        'The function `split_text_into_chunks` should output an array of text chunks.'
                                         'Implement the strategy provided by the user to help split text.')},
         {"role": "user", "content": chunking_strategy}
     ]
@@ -53,13 +53,13 @@ def semantic_chunking_code(text:str) -> str:
     chunking_code_exec = chunking_code.split("```python")[1].split("```")[0]
     return chunking_code_exec
 
-def semantic_chunking(documents:List[Document], chunking_code_exec: str) -> List[Document]:
+def semantic_chunking(documents:List[NeumDocument], chunking_code_exec: str) -> List[NeumDocument]:
     exec(chunking_code_exec, globals())
     result_doc = []
     for doc in documents:
         results = split_text_into_chunks(doc.page_content)
         for result in results:
-            result_doc.append(Document(page_content=result, metadata=doc.metadata))
+            result_doc.append(NeumDocument(id=doc.id, content=result, metadata=doc.metadata))
     return result_doc