hyperforge-nucliadb 1.0.0.post22__py3-none-any.whl

hyperforge_nucliadb/ask/models.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+from hyperforge.models import FieldTypes
+from nuclia.lib.nua_responses import SemanticConfig
+from nuclia_models.worker.triggers import Relation
+from nucliadb_models.graph.responses import (
+    GraphNodesSearchResponse,
+    GraphRelationsSearchResponse,
+    GraphSearchResponse,
+)
+from nucliadb_models.search import Filter, KnowledgeboxFindResults, KnowledgeGraphEntity
+
+
+@dataclass
+class NDBChunk:
+    chunk_id: str
+    text: str
+    page_with_visual: bool
+    reference: Optional[str]
+    start: Optional[int]
+    end: Optional[int]
+    page: Optional[int]
+    labels: List[str]
+    field: Optional[FieldTypes]
+    resource_labels: List[str]
+    link: Optional[str]
+
+
+@dataclass
+class Analysis:
+    semantic_query: str
+    lexical_query: str
+    visual: bool
+    keywords_filter: List[str]
+    reason: Optional[str]
+    input_tokens: Optional[float]
+    output_tokens: Optional[float]
+    labels: List[Filter]
+    semantic_model: Optional[str]
+    entities: List[KnowledgeGraphEntity]
+    pre_queries: List[str]
+    relations: List[Relation]
+    new_text_fields: List[str]
+    new_json_fields: dict[str, Any]
+    semantic_config: SemanticConfig
+    top_k: int
+    link: bool
+    knowledge_graph: Optional[str]
+
+
+@dataclass
+class SearchResults:
+    lexical_find_results: Optional[KnowledgeboxFindResults]
+    semantic_find_results: Optional[KnowledgeboxFindResults]
+    graph_find_results: Optional[KnowledgeboxFindResults]
+    graph_results: Optional[GraphSearchResponse]
+    nodes_results: Optional[GraphNodesSearchResponse]
+    relations_results: Optional[GraphRelationsSearchResponse]

hyperforge_nucliadb/ask/multi.py

@@ -0,0 +1,146 @@
+import asyncio
+from time import time
+from typing import List, cast
+
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import QuestionMemory
+from hyperforge.models import Source
+
+from hyperforge import PROMPT_ENVIRONMENT
+from hyperforge_nucliadb.driver import NucliaDBDriver
+from hyperforge_nucliadb.sync.driver import SyncDriver
+
+MULTI_KB_PROMPT = """
+
+You have multiple Knowledge Boxes (KBs) that you can use as data sources to ask your questions.
+Every KB has an ID, a description, a set of labels, and  facets that characterize its contents.
+Your task is to figure out which KB is the most relevant to a given question by applying the important rules specified.
+
+ Here are the details of each KB:
+
+{% for source in sources.values() %}
+{%- set sid = source.id -%}
+{%- set sdescription = source.description -%}
+{%- set slabels = source.labels -%}
+# ID: {{sid}}
+
+{{sdescription}}
+
+labels: {{slabels}}
+
+{% endfor -%}
+
+And given the question: {{question}}
+
+# Important rules to follow
+
+{% for rule in rules %}
+{{rule}}
+{% endfor -%}
+
+Now, provide your answer following the answer rules:
+- Which KB is the most relevant and why(short explanation)?
+"""
+
+MULTI_KB_PROMPT_TEMPLATE = PROMPT_ENVIRONMENT.from_string(MULTI_KB_PROMPT)
+
+MULTI_KB_SCHEMA = {
+    "title": "source_choice",
+    "description": "Choose which KBs should be used to answer the question of the user and why",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "sources": {
+                "type": "array",
+                "description": "the ids of the choosen sources",
+                "items": {"type": "string"},
+            },
+            "reason": {
+                "type": "string",
+                "description": "reasoning behind the decision",
+            },
+        },
+    },
+}
+
+
+async def choose_source(
+    memory: QuestionMemory,
+    manager: Manager,
+    sources: List[str],
+    question: str,
+    ident: str,
+    choose_source_model: str = "chatgpt-o3-mini",
+    step_title: str = "Knowledge Box Ask: Choose sources",
+) -> list[Source]:
+    if len(sources) == 0:
+        raise Exception("No sources available")
+
+    sources_objects = {}
+    for source_id in sources:
+        source = await memory.get_session_source(source_id)
+        if source is None:
+            source = await load_source_information(source_id, manager)
+            await memory.set_session_source(source)
+        sources_objects[source_id] = source
+
+    # Only one source, try to use it without involving the LLM
+    if len(sources) == 1:
+        return [sources_objects[sources[0]]]
+
+    t0 = time()
+    prompt = MULTI_KB_PROMPT_TEMPLATE.render(
+        sources=sources_objects,
+        rules=memory.get_rules(),
+        question=question,
+    )
+    data, input, output = await manager.execute_json(
+        prompt=prompt,
+        schema=MULTI_KB_SCHEMA,
+        user_id="source_models",
+        model=choose_source_model,
+        tracking=memory.get_tracking_info(),
+    )
+
+    chosen_sources = cast(list[str], data.get("sources"))
+    reason = data.get("reason")
+
+    await memory.add_step(
+        step_module="router",
+        step_title=step_title,
+        step_value=str(chosen_sources),
+        step_reason=reason,
+        timeit=time() - t0,
+        input_nuclia_tokens=input,
+        output_nuclia_tokens=output,
+        step_agent_path=f"/context/{ident if ident else 'default'}",
+    )
+
+    return [sources_objects[chosen] for chosen in chosen_sources]
+
+
+async def load_source_information(source: str, manager: Manager) -> Source:
+    driver = manager.drivers.get(source)
+    if not isinstance(driver, NucliaDBDriver) and not isinstance(driver, SyncDriver):
+        raise ValueError("Source is not a NucliaDB driver")
+    (
+        description,
+        labels,
+        facets_native,
+        paragraph_facets,
+        learning_configuration,
+    ) = await asyncio.gather(
+        driver.description(),
+        driver.labels(),
+        driver.facets_native(),
+        driver.paragraph_facets(),
+        driver.get_learning_configuration(),
+    )
+    return Source(
+        id=source,
+        description=description,
+        labels=labels,
+        facets_native=facets_native,
+        paragraph_facets=paragraph_facets,
+        learning_configuration=learning_configuration,
+    )

hyperforge_nucliadb/ask/nucliadb.py

@@ -0,0 +1,238 @@
+from asyncio import gather
+from typing import Awaitable, List, Optional, cast
+
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import QuestionMemory
+from nucliadb_models.filters import And, Or
+from nucliadb_models.graph.requests import (
+    GraphPathQuery,
+)
+from nucliadb_models.resource import ExtractedDataTypeName
+from nucliadb_models.search import (
+    FindOptions,
+    FindRequest,
+    KnowledgeboxFindResults,
+    MinScore,
+    RerankerName,
+    ResourceProperties,
+)
+from nucliadb_models.security import RequestSecurity
+
+from hyperforge_nucliadb.ask.config import AskAgentConfig
+from hyperforge_nucliadb.ask.models import Analysis, SearchResults
+from hyperforge_nucliadb.ask.utils import (
+    empty,
+    get_nodes,
+    get_relations,
+)
+from hyperforge_nucliadb.driver import NucliaDBDriver
+
+
+def query_relations(
+    memory: QuestionMemory,
+    manager: Manager,
+    analysis: Analysis,
+    nucliadb_driver: NucliaDBDriver,
+    config: AskAgentConfig,
+    relations: List[GraphPathQuery],
+    nodes: List[GraphPathQuery],
+) -> Awaitable[KnowledgeboxFindResults]:
+    """Query relations in NucliaDB based on the provided analysis."""
+    find_request = FindRequest(
+        features=[FindOptions.GRAPH],
+        filters=analysis.labels,
+        graph_query=And(
+            operands=[
+                Or(operands=relations),
+                Or(operands=nodes),
+            ]
+        ),
+        security=RequestSecurity(groups=config.security_groups),
+        reranker=RerankerName.NOOP,
+        show=[
+            ResourceProperties.BASIC,
+            ResourceProperties.ORIGIN,
+            ResourceProperties.EXTRA,
+            ResourceProperties.VALUES,
+            ResourceProperties.RELATIONS,
+        ],
+    )
+    return nucliadb_driver.find_raw(find_request)
+
+
+def query_semantic(
+    memory: QuestionMemory,
+    manager: Manager,
+    analysis: Analysis,
+    nucliadb_driver: NucliaDBDriver,
+    config: AskAgentConfig,
+) -> Awaitable[KnowledgeboxFindResults]:
+    """Semantic search in NucliaDB based on the provided analysis."""
+
+    find_request = FindRequest(
+        features=[FindOptions.SEMANTIC],
+        query=analysis.semantic_query,
+        min_score=MinScore(semantic=analysis.semantic_config.threshold),
+        vectorset=analysis.semantic_model,
+        filters=analysis.labels,
+        security=RequestSecurity(groups=config.security_groups),
+        reranker=RerankerName.NOOP,
+        show=[
+            ResourceProperties.BASIC,
+            ResourceProperties.ORIGIN,
+            ResourceProperties.EXTRA,
+            ResourceProperties.EXTRACTED,
+            ResourceProperties.VALUES,
+            ResourceProperties.RELATIONS,
+        ],
+        extracted=[
+            ExtractedDataTypeName.TEXT,
+            ExtractedDataTypeName.METADATA,
+            ExtractedDataTypeName.FILE,
+            ExtractedDataTypeName.LINK,
+        ],
+    )
+
+    return nucliadb_driver.find_raw(find_request)
+
+
+def query_lexical(
+    memory: QuestionMemory,
+    manager: Manager,
+    analysis: Analysis,
+    nucliadb_driver: NucliaDBDriver,
+    config: AskAgentConfig,
+) -> Awaitable[KnowledgeboxFindResults]:
+    """Lexical search (aka keyword search) in NucliaDB based on the provided analysis."""
+
+    find_request = FindRequest(
+        features=[FindOptions.KEYWORD],
+        query=analysis.lexical_query,
+        filters=analysis.labels,
+        keyword_filters=analysis.keywords_filter,
+        security=RequestSecurity(groups=config.security_groups),
+        reranker=RerankerName.NOOP,
+        show=[
+            ResourceProperties.BASIC,
+            ResourceProperties.ORIGIN,
+            ResourceProperties.EXTRA,
+            ResourceProperties.EXTRACTED,
+            ResourceProperties.VALUES,
+            ResourceProperties.RELATIONS,
+        ],
+        extracted=[
+            ExtractedDataTypeName.TEXT,
+            ExtractedDataTypeName.METADATA,
+            ExtractedDataTypeName.FILE,
+            ExtractedDataTypeName.LINK,
+        ],
+    )
+
+    return nucliadb_driver.find_raw(find_request)
+
+
+async def query_ndb(
+    memory: QuestionMemory,
+    manager: Manager,
+    analysis: Analysis,
+    nucliadb_driver: NucliaDBDriver,
+    config: AskAgentConfig,
+) -> SearchResults:
+    semantic_paragraphs = query_semantic(
+        memory, manager, analysis, nucliadb_driver, config
+    )
+    lexical_paragraphs = query_lexical(
+        memory, manager, analysis, nucliadb_driver, config
+    )
+    nodes = cast(list[GraphPathQuery], get_nodes(analysis))
+    relations = cast(list[GraphPathQuery], get_relations(analysis))
+
+    if len(relations) > 0 and len(nodes) > 0:
+        relations_paragraphs: Awaitable[Optional[KnowledgeboxFindResults]] = (
+            query_relations(
+                memory,
+                manager,
+                analysis,
+                nucliadb_driver,
+                config,
+                relations=cast(list[GraphPathQuery], relations),
+                nodes=cast(list[GraphPathQuery], nodes),
+            )
+        )
+
+    else:
+        relations_paragraphs = empty()
+
+    (
+        lexical_find_results,
+        semantic_find_results,
+        graph_find_results,
+    ) = await gather(
+        lexical_paragraphs,
+        semantic_paragraphs,
+        relations_paragraphs,
+    )
+
+    if (
+        len(lexical_find_results.resources) == 0
+        and analysis.keywords_filter is not None
+        and len(analysis.keywords_filter) > 0
+    ):
+        # XXX: This is repeating the lexical query with the same parameters no?
+
+        # If no lexical results, we can try to find some using the keywords filter
+        lexical_find_results = await query_lexical(
+            memory, manager, analysis, nucliadb_driver, config
+        )
+
+    result = SearchResults(
+        lexical_find_results=lexical_find_results,
+        semantic_find_results=semantic_find_results,
+        graph_find_results=graph_find_results,
+        graph_results=None,
+        nodes_results=None,
+        relations_results=None,
+    )
+    return result
+
+
+async def standard_query_ndb(
+    memory: QuestionMemory,
+    manager: Manager,
+    nucliadb_driver: NucliaDBDriver,
+    config: AskAgentConfig,
+    question: str,
+) -> SearchResults:
+    find_request = FindRequest(
+        features=[FindOptions.SEMANTIC, FindOptions.KEYWORD],
+        query=question,
+        security=RequestSecurity(groups=config.security_groups),
+        reranker=RerankerName.NOOP,
+        show=[
+            ResourceProperties.BASIC,
+            ResourceProperties.ORIGIN,
+            ResourceProperties.EXTRA,
+            ResourceProperties.EXTRACTED,
+            ResourceProperties.VALUES,
+            ResourceProperties.RELATIONS,
+        ],
+        extracted=[
+            ExtractedDataTypeName.TEXT,
+            ExtractedDataTypeName.METADATA,
+            ExtractedDataTypeName.FILE,
+            ExtractedDataTypeName.LINK,
+        ],
+        filters=nucliadb_driver.config.filters,
+    )
+
+    paragraphs = await nucliadb_driver.find_raw(find_request)
+
+    result = SearchResults(
+        lexical_find_results=None,
+        semantic_find_results=paragraphs,
+        graph_find_results=None,
+        graph_results=None,
+        nodes_results=None,
+        relations_results=None,
+    )
+    return result

hyperforge_nucliadb/ask/prompt_analysis.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict
+
+from hyperforge_nucliadb.ask.config import AskAgentConfig
+from hyperforge_nucliadb.ask.kb_analysis import KnowledgeBoxInfo
+from hyperforge_nucliadb.ask.query_analysis import QueryAnalisys
+
+# https://huggingface.co/collections/ModelSpace/gemmax2-673714f5049bfa3a90bee6b6
+TRANSLATE_QUERY_PROMPT = "Rephrase this question so its better for lexical retrieval, translate the lexical rephrased question to {language}. Please define ONLY the question without any explanation. JUST A SENTENCE. DO NOT ADD ANY EXTRA NOTES AT THE END, JUST ONE SENTENCE"
+
+
+def configure_prompts(
+    config: AskAgentConfig,
+    ask_json_schema_copy: Dict[str, Dict[str, Any]],
+    kb_analysis: KnowledgeBoxInfo,
+    query_analysis: QueryAnalisys,
+) -> None:
+    """
+    Configure the prompts based on the provided configuration.
+    This function modifies the ask_json_schema_copy in place to set the
+    appropriate prompts for semantic, lexical, and visual queries.
+    """
+    if config.visual_enable_prompt is not None:
+        ask_json_schema_copy["parameters"]["properties"]["visual"]["description"] = (
+            config.visual_enable_prompt
+        )
+
+    if config.rephrase_semantic_custom_prompt is not None:
+        ask_json_schema_copy["parameters"]["properties"]["semantic_query"][
+            "description"
+        ] = config.rephrase_semantic_custom_prompt
+
+    # Rephrase as lexical to query lexical engine
+    if config.rephrase_lexical_custom_prompt is not None:
+        ask_json_schema_copy["parameters"]["properties"]["lexical_query"][
+            "description"
+        ] = config.rephrase_lexical_custom_prompt
+    elif query_analysis.translate_need is not None:
+        ask_json_schema_copy["parameters"]["properties"]["lexical_query"][
+            "description"
+        ] = TRANSLATE_QUERY_PROMPT.format(language=query_analysis.translate_need)
+
+    # Rephrase as keywords that must appear only if the language of the query matches the language of the KB
+
+    if (
+        query_analysis.translate_need is None
+        and config.keywords_custom_prompt is not None
+    ):
+        ask_json_schema_copy["parameters"]["properties"]["keywords_filter"][
+            "description"
+        ] = config.keywords_custom_prompt

hyperforge_nucliadb/ask/query_analysis.py

@@ -0,0 +1,83 @@
+import re
+from dataclasses import dataclass
+from typing import Optional
+
+from lingua import LanguageDetectorBuilder
+
+from hyperforge import logger
+from hyperforge_nucliadb.ask.config import AskAgentConfig
+from hyperforge_nucliadb.ask.kb_analysis import KnowledgeBoxInfo
+
+LID_MODEL = None
+
+
+def load_lingua() -> None:
+    global LID_MODEL
+    LID_MODEL = (
+        LanguageDetectorBuilder.from_all_languages()
+        .with_preloaded_language_models()
+        .build()
+    )
+
+
+def detect_language(text: str, default="en") -> str:
+    global LID_MODEL
+    lang = default
+
+    if len(text.strip()) == 0:
+        return default
+
+    if isinstance(text, bytes):
+        text = text.decode()
+
+    text = cleanup_text_for_langdetect(text)
+
+    if LID_MODEL is None:
+        load_lingua()
+    assert LID_MODEL is not None  # for type checker
+    result = LID_MODEL.detect_language_of(text.replace("\n", " "))
+
+    if result is not None:
+        lang = result.iso_code_639_1.name.lower()
+
+    return lang
+
+
+def cleanup_text_for_langdetect(text: str) -> str:
+    # Get all text from splitting docs
+    regex = re.compile(r"(\[image: [a-zA-Z0-9.]*\])", re.S)
+    return re.sub(regex, "", text)
+
+
+@dataclass
+class QueryAnalisys:
+    translate_need: Optional[str] = None
+
+
+def pre_query_analysis(
+    config: AskAgentConfig, question: str, kb_analysis: KnowledgeBoxInfo
+) -> QueryAnalisys:
+    # LLM QUERY
+    result = QueryAnalisys()
+
+    # Rephrase as semantic to query semantic engine -> custom prompt
+
+    query_language = detect_language(question)
+
+    # Check if the query language is in the languages of the KB
+    # or get the most common language in the KB to translate the query
+    if query_language in kb_analysis.languages:
+        result.translate_need = None
+    else:
+        try:
+            result.translate_need = sorted(
+                kb_analysis.languages.items(), key=lambda item: item[1]
+            )[0][0]
+        except IndexError:
+            logger.warning(
+                "No languages found in the KB analysis, cannot determine if translation is needed"
+            )
+            # Kb may not have any languages
+            result.translate_need = None
+
+    return result

hyperforge_nucliadb/ask/rerank.py

@@ -0,0 +1,45 @@
+from hyperforge.manager import Manager
+from hyperforge.models import Context
+from nuclia.lib.nua_responses import RerankModel
+
+from hyperforge import logger
+
+
+async def rerank(context: Context, manager: Manager, top_k: int) -> None:
+    """
+    Rerank the context chunks based on the analysis.
+    """
+    if not context.chunks:
+        return
+
+    re_ranked = await manager.rerank(
+        RerankModel(
+            context={chunk.chunk_id: chunk.text for chunk in context.chunks},
+            question=context.question,
+            user_id="arag-ask-rerank",
+        )
+    )
+
+    # Sort chunks by score
+    re_ranked_sorted = sorted(
+        re_ranked.context_scores.items(), key=lambda x: x[1], reverse=True
+    )
+
+    re_ranked_sorted = list(
+        filter(
+            lambda x: x[1] > 0.05,
+            re_ranked.context_scores.items(),
+        )
+    )
+
+    re_ranked_sorted = re_ranked_sorted[:top_k]
+
+    # If the analysis has a rerank function, use it to rerank the chunks
+    old_chunks = {chunk.chunk_id: chunk for chunk in context.chunks}
+    context.chunks = []
+    for rerank_chunk_id, _ in re_ranked_sorted:
+        try:
+            chunk = old_chunks[rerank_chunk_id]
+            context.chunks.append(chunk)
+        except KeyError:
+            logger.error(f"Chunk {rerank_chunk_id} not found in old_chunks")

hyperforge_nucliadb/ask/utils.py

@@ -0,0 +1,36 @@
+from collections.abc import Sequence
+from typing import List
+
+from nucliadb_models.graph.requests import (
+    AnyNode,
+    NodeMatchKindName,
+    Relation,
+)
+
+from hyperforge_nucliadb.ask.models import Analysis
+
+
+async def empty() -> None:
+    """A placeholder to call."""
+    pass
+
+
+def get_nodes(analysis: Analysis) -> List[AnyNode]:
+    """Extract nodes from the analysis."""
+    nodes: List[AnyNode] = []
+    for entity in analysis.entities:
+        if entity.type == "PERSON":
+            nodes.append(AnyNode(value=entity.name, match=NodeMatchKindName.FUZZY))
+        elif entity.type == "LOCATION" or entity.type == "ORGANIZATION":
+            nodes.append(AnyNode(value=entity.name, match=NodeMatchKindName.EXACT))
+        else:
+            nodes.append(AnyNode(value=entity.name, match=NodeMatchKindName.FUZZY))
+    return nodes
+
+
+def get_relations(analysis: Analysis) -> Sequence[Relation]:
+    """Extract relations from the analysis."""
+    relations: List[Relation] = []
+    for relation in analysis.relations:
+        relations.append(Relation(label=relation.label))
+    return relations