veri-agents-knowledgebase 0.1.1__py3-none-any.whl

veri_agents_knowledgebase/__init__.py

@@ -0,0 +1 @@
+from .knowledgebase import DocumentLoader, KnowledgeFilter, DataSource, KnowledgebaseMetadata, Knowledgebase, RWKnowledgebase

veri_agents_knowledgebase/graphs/__init__.py

@@ -0,0 +1,2 @@
+from .qa import create_qa_agent
+from .extract import create_extract_agent

veri_agents_knowledgebase/graphs/extract.py

@@ -0,0 +1,213 @@
+import logging
+from typing import Annotated, Callable, List, Optional, Sequence, TypedDict, cast
+
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.graph import CompiledGraph
+from langgraph.graph.message import add_messages
+from pydantic import BaseModel, Field
+from veri_agents_knowledgebase import Knowledgebase, KnowledgeFilter
+from veri_agents_knowledgebase.utils import get_filter_from_config, aprocess_docs
+
+log = logging.getLogger(__name__)
+log.setLevel(logging.DEBUG)
+
+Entities = dict[str, list[str]]
+EntitiesPerDocument = dict[str, Entities]
+
+
+class ExtractionResponse(BaseModel):
+    """Extracted information from documents"""
+
+    entities: Entities = Field(
+        description="Entities extracted from the documents. If you can't find any occurrences of the requested entity, return an empty list for that entity. DO NOT call this tool with the provided examples. Don't return multiple variants of the same entity like 'employee' and 'employees'. Type of field is JSON object with keys corresponding to entity type, and values of entity name with type string. NOT A STRING",
+        examples=[
+            {
+                "persons": [
+                    "John Doe",
+                    "Jane James",
+                    "S1 (Construction Superintendent)",
+                ],
+                "objects": ["excavator", "dog"],
+            },
+            {
+                "locations": ["New York", "Los Angeles"],
+                "dates": ["2025-02-05", "2025-03-10"],
+            },
+        ],
+    )
+
+
+def reduce_extraction_responses(
+    existing: Entities | None, new: ExtractionResponse
+) -> Entities:
+    """Reduce two extraction responses by merging their entities."""
+    if existing is None:
+        existing = Entities()
+    for entity_type, entities in new.entities.items():
+        if entity_type not in existing:
+            existing[entity_type] = []
+        # Add only new entities (avoid duplicates)
+        for entity in entities:
+            if entity not in existing[entity_type]:
+                existing[entity_type].append(entity)
+    return existing
+
+
+class ExtractInputSchema(BaseModel):
+    """Input schema for the extract workflow"""
+
+    knowledgebase: Optional[str] = Field(
+        default=None,
+        description="Knowledgebase to extract from",
+        examples=["knowledgebase1", "knowledgebase2"],
+    )
+    documents: List[str] = Field(
+        default=[],
+        description="Extract just within the list of documents with the given IDs.",
+        examples=[["source1", "source2"]],
+    )
+    tags_any: List[str] = Field(
+        default=[],
+        description="Extract entities from documents matching any of the given tags.",
+        examples=[["finance", "priority_high"]],
+    )
+    tags_all: List[str] = Field(
+        default=[],
+        description="Extract entities from documents matching all of the given tags.",
+        examples=[["finance", "priority_high"]],
+    )
+    llm: Optional[str] = Field(
+        default=None,
+        description="LLM to use for extraction",
+        examples=["nova_pro", "bedrock_claude_sonnet_37"],
+    )
+
+
+class ExtractOutputSchema(BaseModel):
+    """Output schema for the extract workflow"""
+
+    content: str = Field(
+        description="Short summary of the entities",
+        examples=[
+            "The following entities were extracted: John Doe, Jane James, excavator"
+        ],
+    )
+    entities: EntitiesPerDocument = Field(
+        default={},
+        description="Extracted entities. Dictionary of documents to entities where entities are a dictionary of entity types to list of entities",
+    )
+
+
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[AnyMessage], add_messages]
+    documents: EntitiesPerDocument
+
+
+def create_extract_agent(
+    extract_llm: BaseLanguageModel | Callable[[RunnableConfig], BaseLanguageModel],
+    knowledgebase: Knowledgebase | Callable[[RunnableConfig], Knowledgebase],
+    filter: KnowledgeFilter | None = None,
+    system_prompt: str | None = None,
+    summarize_prompt: str | None = None,
+) -> CompiledGraph:
+    if system_prompt is None:
+        system_prompt = "Your job is to identify important entities, aspects, key points, incidents etc. in the provided data. only use the provided context, don't use your internal knowledge to make up things."
+    if summarize_prompt is None:
+        summarize_prompt = "Summarize your findings about the extracted entities from the documents: \n"
+
+    async def aextract(state: AgentState, config: RunnableConfig) -> AgentState:
+        """LangGraph node to extract entities from the documents"""
+        try:
+            messages = state["messages"]
+            runnable_filter = get_filter_from_config(config, filter, prefix="filter_")
+            print("EXTRACT RUNNABLE FILTER")
+            print(runnable_filter)
+            runnable_kb = knowledgebase(config) if callable(knowledgebase) else knowledgebase
+            runnable_llm = extract_llm(config) if not isinstance(extract_llm, BaseLanguageModel) else extract_llm
+            docs = runnable_kb.get_documents(runnable_filter)
+            llm_structured = runnable_llm.with_structured_output(ExtractionResponse)
+
+            extraction_results = await aprocess_docs(
+                docs,
+                messages,
+                llm_structured,
+                system_prompt=system_prompt,
+                reduce_results=reduce_extraction_responses,
+            )
+            extraction_results = cast(EntitiesPerDocument, extraction_results)
+            artifact = {
+                "type": "json",
+                "source": "extracted_entities",
+                "documents": extraction_results,
+            }
+            return {
+                "messages": [
+                    ToolMessage(
+                        "Extraction complete.",
+                        tool_call_id="extract",
+                        artifact=artifact,
+                    )
+                ],
+                "documents": extraction_results,
+            }
+        except Exception as e:
+            log.error("Error in extract: ", exc_info=True)
+            return {
+                "messages": [AIMessage(content="Error in extraction: " + str(e))],
+                "documents": {},
+            }
+
+    async def asummarize(state: AgentState, config: RunnableConfig) -> AgentState:
+        """LangGraph node to summarize the extracted entities from previous nodes."""
+        runnable_llm = extract_llm(config) if not isinstance(extract_llm, BaseLanguageModel) else extract_llm
+        try:
+            messages = state["messages"]
+            send_messages = [
+                SystemMessage(content=system_prompt),
+                messages[-2],  # the user input
+                HumanMessage(content=summarize_prompt + str(state["documents"])),
+            ]
+
+            response = await runnable_llm.ainvoke(send_messages)
+            output = ExtractOutputSchema(
+                content=response.content,
+                entities=state["documents"],
+            )
+
+            return {
+                "messages": [
+                    AIMessage(
+                        content=output.model_dump_json(),
+                    )
+                ],
+                "documents": state["documents"],
+            }
+        except Exception as e:
+            log.error("Error in summarize: ", exc_info=True)
+            output = ExtractOutputSchema(
+                content="Error in summarization: " + str(e),
+                entities=state["documents"],
+            )
+            return {
+                "messages": [AIMessage(content=output.model_dump_json())],
+                "documents": state["documents"],
+            }
+
+    graph = StateGraph(AgentState)
+    graph.add_node("extract", aextract)
+    graph.add_node("summarize", asummarize)
+
+    graph.add_edge(START, "extract")
+    graph.add_edge("extract", "summarize")
+    graph.add_edge("summarize", END)
+
+    return graph.compile()

veri_agents_knowledgebase/graphs/qa.py

@@ -0,0 +1,51 @@
+import logging
+from datetime import datetime
+from typing import Sequence, Callable
+
+from langchain_core.language_models import (
+    LanguageModelLike,
+)
+from langgraph.graph.graph import CompiledGraph
+from langgraph.prebuilt import ToolNode, create_react_agent
+from langchain_core.tools import BaseTool
+
+from veri_agents_knowledgebase import Knowledgebase
+from veri_agents_knowledgebase.tools import FixedKnowledgebaseWithTagsQuery, FixedKnowledgebaseListDocuments
+
+log = logging.getLogger(__name__)
+
+def create_qa_agent(
+    llm: LanguageModelLike,
+    knowledgebases: Sequence[Knowledgebase],
+    system_prompt: str,
+    tools: Sequence[BaseTool | Callable] | None = None,
+    **react_kwargs
+) -> CompiledGraph:
+    tools = list(tools) if tools else []
+    for i, knowledgebase in enumerate(knowledgebases):
+        tools.append(
+            FixedKnowledgebaseWithTagsQuery(
+                knowledgebase=knowledgebase,
+                num_results=10,
+                name_suffix=f"-{i}",
+                runnable_config_filter_prefix="filter_",  # TODO: pick your own prefix for the runnable config
+            )
+        )
+        tools.append(
+           FixedKnowledgebaseListDocuments(
+               knowledgebase=knowledgebase,
+               name_suffix=f"-{i}",
+               runnable_config_filter_prefix="filter_",  # TODO: pick your own prefix for the runnable config
+           ) 
+        )
+    tool_node = ToolNode(tools)
+
+    system_prompt = system_prompt
+    system_prompt += f"""Today's date is: {datetime.now().strftime("%Y-%m-%d")}."""
+
+    return create_react_agent(
+        model=llm,
+        tools=tool_node,
+        prompt=system_prompt,
+        **react_kwargs
+    )

veri_agents_knowledgebase/knowledgebase.py

@@ -0,0 +1,178 @@
+import logging
+from os import PathLike
+from typing import Iterator, Optional, List, Dict, cast
+from collections.abc import Sequence
+
+from langchain_core.documents import Document
+from pydantic import BaseModel, Field
+
+log = logging.getLogger(__name__)
+
+
+class DocumentLoader:
+    """Loads data from a data source and returns documents."""
+
+    def __init__(self):
+        pass
+
+    def load_documents(self, **kwargs):
+        """Parse documents from a data source."""
+        raise NotImplementedError
+
+
+class KnowledgeFilter(BaseModel):
+    """Filter for knowledge base queries."""
+
+    docs: list[str] | str | None = None
+    """ List of document IDs or single document ID to filter by. """
+
+    tags_any_of: list[str] | str | None = None
+    """ List of tags to filter by, if any of the provided tags matches, a document is selected. """
+
+    tags_all_of: list[str] | str | None = None
+    """ List of tags to filter by, if all of the provided tags match, a document is selected. """
+
+    def __repr__(self):
+        return f"KnowledgeFilter(docs={self.docs}, tags_any_of={self.tags_any_of}, tags_all_of={self.tags_all_of})"
+
+    def __str__(self):
+        return f"KnowledgeFilter(docs={self.docs}, tags_any_of={self.tags_any_of}, tags_all_of={self.tags_all_of})"
+
+def and_filters(filter1: KnowledgeFilter | None, filter2: KnowledgeFilter | None):
+    if filter1 is None:
+        return filter2
+    elif filter2 is None:
+        return filter1
+    else:
+        # docs
+        if filter1.docs is None:
+            docs = filter2.docs
+        elif filter2.docs is None:
+            docs = filter1.docs
+        else:
+            # intersection
+            docs1 = filter1.docs if isinstance(filter1.docs, Sequence) and not isinstance(filter1.docs, str) else [cast(str, filter1.docs)]
+            docs2 = filter2.docs if isinstance(filter2.docs, Sequence) and not isinstance(filter2.docs, str) else [cast(str, filter2.docs)]
+
+            docs = list(set(docs1) & set(docs2))
+
+        # tags_any_of
+        if filter1.tags_any_of is None:
+            tags_any_of = filter2.tags_any_of
+        elif filter2.tags_any_of is None:
+            tags_any_of = filter1.tags_any_of
+        else:
+            # union
+            tags_any_of1 = filter1.tags_any_of if isinstance(filter1.tags_any_of, Sequence) and not isinstance(filter1.tags_any_of, str) else [cast(str, filter1.tags_any_of)]
+            tags_any_of2 = filter2.tags_any_of if isinstance(filter2.tags_any_of, Sequence) and not isinstance(filter2.tags_any_of, str) else [cast(str, filter2.tags_any_of)]
+
+            tags_any_of = list(set(tags_any_of1) | set(tags_any_of2))
+
+        # tags_all_of
+        if filter1.tags_all_of is None:
+            tags_all_of = filter2.tags_all_of
+        elif filter2.tags_all_of is None:
+            tags_all_of = filter1.tags_all_of
+        else:
+            # union
+            tags_all_of1 = filter1.tags_all_of if isinstance(filter1.tags_all_of, Sequence) and not isinstance(filter1.tags_all_of, str) else [cast(str, filter1.tags_all_of)]
+            tags_all_of2 = filter2.tags_all_of if isinstance(filter2.tags_all_of, Sequence) and not isinstance(filter2.tags_all_of, str) else [cast(str, filter2.tags_all_of)]
+
+            tags_all_of = list(set(tags_all_of1) | set(tags_all_of2))
+
+        return KnowledgeFilter(
+            docs=docs,
+            tags_any_of=tags_any_of,
+            tags_all_of=tags_all_of
+        )
+
+
+class DataSource(BaseModel):
+    """Data source for a knowledge base."""
+
+    location: PathLike | str = Field(
+        description="Location of the data source, e.g. a file path or URL."
+    )
+    name: str = Field(
+        description="Name of the data source. Can be used for filtering in the knowledgebase and important that document names are unique"
+    )
+    tags: list[str] = Field(
+        default=[],
+        description="Tags applied to all documents and chunks of the source, e.g. 'finance'.",
+    )
+    incremental: bool = Field(
+        default=False,
+        description="Whether to do incremental indexing of the data source.",
+    )
+
+class KnowledgebaseMetadata(BaseModel):
+    """Metadata for a knowledgebase."""
+
+    name: str
+    description: str | None = None
+    tags: dict[str, str] = {}
+    collection: str
+    doc_summarize: bool = False
+    doc_autotag: bool = False
+    data_sources: List[Dict] = []
+
+    class Config:
+        extra = 'ignore'
+
+class Knowledgebase:
+    def __init__(self, **kwargs):
+        print(f"Knowledgebase init: {kwargs}", flush=True)
+        self.metadata = KnowledgebaseMetadata.model_validate(kwargs)
+
+    @property
+    def tags(self):
+        """Get the tags for the workflow."""
+        return self.metadata.tags
+
+    @property
+    def name(self):
+        """Get the name of the workflow."""
+        return self.metadata.name
+
+    @property
+    def description(self):
+        """Get the description of the workflow."""
+        return self.metadata.description
+
+    def retrieve(
+        self,
+        query: str,
+        limit: int,
+        filter: KnowledgeFilter | None = None,
+        **kwargs,
+    ):
+        """Retrieve documents from the knowledge base."""
+        raise NotImplementedError
+
+    def get_documents(
+        self,
+        filter: KnowledgeFilter | None = None,
+    ) -> Iterator[Document]:
+        """Get all documents from the knowledge base."""
+        raise NotImplementedError
+
+
+class RWKnowledgebase(Knowledgebase):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def index(self, data_source: Optional[DataSource] = None):
+        """Do an index run on either a provides data source or data sources defined in its config.
+
+        Args:
+            data_source (DataSource): Data source to index. If None, will use the data sources defined in the config.
+        """
+        raise NotImplementedError
+
+    def set_tags(
+        self,
+        doc_id: str,
+        tags: list[str],
+    ):
+        """Add tags to a document."""
+        raise NotImplementedError

veri_agents_knowledgebase/qdrant/__init__.py

@@ -0,0 +1,5 @@
+from .qdrant_doc_store import QdrantDocStore
+from .qdrant_kb import QdrantKnowledgebase
+from .generic_kb import GenericQdrantKnowledgebase
+from .source_retriever import *
+from .summarization import *