hyperforge 1.0.0.post19__py3-none-any.whl

hyperforge/nua.py

@@ -0,0 +1,336 @@
+from enum import Enum
+from typing import (
+    Any,
+    AsyncIterator,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
+
+import backoff
+from deprecated import deprecated
+from httpx import AsyncClient
+from nuclia.exceptions import NuaAPIException
+from nuclia.lib.nua_responses import (
+    ChatModel,
+    ChatResponse,
+    QueryInfo,
+    RephraseModel,
+    RerankModel,
+    RerankResponse,
+    Sentence,
+    SummarizedModel,
+    SummarizeModel,
+    SummarizeResource,
+    Tokens,
+)
+from nuclia_models.common.consumption import Consumption, ConsumptionGenerative
+from nuclia_models.predict.generative_responses import (
+    CitationsGenerativeResponse,
+    GenerativeChunk,
+    GenerativeFullResponse,
+    JSONGenerativeResponse,
+    MetaGenerativeResponse,
+    StatusGenerativeResponse,
+    TextGenerativeResponse,
+    ToolsGenerativeResponse,
+)
+from nuclia_models.predict.remi import RemiRequest, RemiResponse
+from pydantic import BaseModel
+
+MB = 1024 * 1024
+CHUNK_SIZE = 10 * MB
+SENTENCE_PREDICT = "/api/v1/internal/predict/sentence"
+CHAT_PREDICT = "/api/v1/internal/predict/chat"
+SUMMARIZE_PREDICT = "/api/v1/internal/predict/summarize"
+REPHRASE_PREDICT = "/api/v1/internal/predict/rephrase"
+TOKENS_PREDICT = "/api/v1/internal/predict/tokens"
+QUERY_PREDICT = "/api/v1/internal/predict/query"
+REMI_PREDICT = "/api/v1/internal/predict/remi"
+AGENTS_PREDICT = "/api/v1/internal/predict/run-agents"
+RERANK = "/api/v1/internal/predict/rerank"
+
+ConvertType = TypeVar("ConvertType", bound=BaseModel)
+
+
+class Author(str, Enum):
+    NUCLIA = "NUCLIA"
+    USER = "USER"
+
+
+class ContextItem(BaseModel):
+    author: Author
+    text: str
+
+
+class RetriableRequestException(NuaAPIException):
+    pass
+
+
+class AsyncInternalNuaClient:
+    def __init__(
+        self,
+        kbid: str | None,
+        account: str | None,
+        url: str,
+    ):
+        self.headers = {"X-STF-KBID": kbid} if kbid else {}
+        if account:
+            self.headers["X-STF-ACCOUNT"] = account
+
+        self.stream_headers = self.headers.copy()
+        self.stream_headers["Accept"] = "application/x-ndjson"
+
+        self.url = url
+        self.client = AsyncClient(headers=self.headers, base_url=url)
+        self.stream_client = AsyncClient(headers=self.stream_headers, base_url=url)
+
+    @backoff.on_exception(
+        backoff.expo,
+        (RetriableRequestException,),
+        max_time=60,
+        jitter=backoff.full_jitter,
+    )
+    async def _request(
+        self,
+        method: str,
+        url: str,
+        output: Type[ConvertType],
+        payload: Optional[dict[Any, Any]] = None,
+        timeout: int = 60,
+    ) -> ConvertType:
+        resp = await self.client.request(method, url, json=payload, timeout=timeout)
+        if resp.status_code in (429, 512):
+            raise RetriableRequestException(
+                code=resp.status_code, detail=resp.content.decode()
+            )
+        if resp.status_code > 299:
+            raise NuaAPIException(code=resp.status_code, detail=resp.content.decode())
+        try:
+            data = output.model_validate(resp.json())
+        except Exception:
+            data = output.model_validate(resp.content)
+        return data
+
+    @backoff.on_exception(
+        backoff.expo,
+        (RetriableRequestException,),
+        max_time=60,
+        jitter=backoff.full_jitter,
+    )
+    async def _stream(
+        self,
+        method: str,
+        url: str,
+        payload: Optional[dict[Any, Any]] = None,
+        timeout: int = 60,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> AsyncIterator[GenerativeChunk]:
+        async with self.stream_client.stream(
+            method,
+            url,
+            json=payload,
+            timeout=timeout,
+            headers=extra_headers,
+        ) as response:
+            if response.status_code in (429, 512):
+                raise RetriableRequestException(
+                    code=response.status_code,
+                    detail=(await response.aread()).decode(errors="ignore"),
+                )
+            elif response.status_code > 299:
+                raise NuaAPIException(
+                    code=response.status_code,
+                    detail=(await response.aread()).decode(errors="ignore"),
+                )
+            async for json_body in response.aiter_lines():
+                yield GenerativeChunk.model_validate_json(json_body)
+
+    async def sentence_predict(
+        self, text: str, model: Optional[str] = None
+    ) -> Sentence:
+        endpoint = f"{self.url}{SENTENCE_PREDICT}?text={text}"
+        if model:
+            endpoint += f"&model={model}"
+        return await self._request("GET", endpoint, output=Sentence)
+
+    async def tokens_predict(self, text: str, model: Optional[str] = None) -> Tokens:
+        endpoint = f"{self.url}{TOKENS_PREDICT}?text={text}"
+        if model:
+            endpoint += f"&model={model}"
+        return await self._request("GET", endpoint, output=Tokens)
+
+    async def query_predict(
+        self,
+        text: str,
+        semantic_model: Optional[str] = None,
+        token_model: Optional[str] = None,
+        generative_model: Optional[str] = None,
+    ) -> QueryInfo:
+        endpoint = f"{self.url}{QUERY_PREDICT}?text={text}"
+        if semantic_model:
+            endpoint += f"&semantic_model={semantic_model}"
+        if token_model:
+            endpoint += f"&token_model={token_model}"
+        if generative_model:
+            endpoint += f"&generative_model={generative_model}"
+        return await self._request("GET", endpoint, output=QueryInfo)
+
+    @deprecated(version="2.1.0", reason="You should use generate function")
+    async def generate_predict(
+        self, body: ChatModel, model: Optional[str] = None, timeout: int = 300
+    ) -> ChatResponse:
+        endpoint = f"{self.url}{CHAT_PREDICT}"
+        if model:
+            endpoint += f"?model={model}"
+
+        return await self._request(
+            "POST",
+            endpoint,
+            payload=body.model_dump(),
+            output=ChatResponse,
+            timeout=timeout,
+        )
+
+    async def generate(
+        self,
+        body: ChatModel,
+        model: Optional[str] = None,
+        timeout: int = 300,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> GenerativeFullResponse:
+        endpoint = f"{self.url}{CHAT_PREDICT}"
+        if model:
+            endpoint += f"?model={model}"
+        result = GenerativeFullResponse(answer="")
+        async for chunk in self._stream(
+            "POST",
+            endpoint,
+            payload=body.model_dump(),
+            timeout=timeout,
+            extra_headers=extra_headers,
+        ):
+            if isinstance(chunk.chunk, TextGenerativeResponse):
+                result.answer += chunk.chunk.text
+            elif isinstance(chunk.chunk, JSONGenerativeResponse):
+                result.object = chunk.chunk.object
+            elif isinstance(chunk.chunk, MetaGenerativeResponse):
+                result.timings = chunk.chunk.timings
+            elif isinstance(chunk.chunk, CitationsGenerativeResponse):
+                result.citations = chunk.chunk.citations
+            elif isinstance(chunk.chunk, StatusGenerativeResponse):
+                result.code = chunk.chunk.code
+            elif isinstance(chunk.chunk, ToolsGenerativeResponse):
+                result.tools = chunk.chunk.tools
+            elif isinstance(chunk.chunk, ConsumptionGenerative):
+                result.consumption = Consumption(
+                    normalized_tokens=chunk.chunk.normalized_tokens,
+                    customer_key_tokens=chunk.chunk.customer_key_tokens,
+                )
+        return result
+
+    async def generate_stream(
+        self,
+        body: ChatModel,
+        model: Optional[str] = None,
+        timeout: int = 300,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> AsyncIterator[GenerativeChunk]:
+        endpoint = f"{self.url}{CHAT_PREDICT}"
+        if model:
+            endpoint += f"?model={model}"
+
+        async for gr in self._stream(
+            "POST",
+            endpoint,
+            payload=body.model_dump(),
+            timeout=timeout,
+        ):
+            yield gr
+
+    async def summarize(
+        self, documents: dict[str, str], model: Optional[str] = None, timeout: int = 300
+    ) -> SummarizedModel:
+        endpoint = f"{self.url}{SUMMARIZE_PREDICT}"
+        if model:
+            endpoint += f"?model={model}"
+
+        body = SummarizeModel(
+            resources={
+                key: SummarizeResource(fields={"field": document})
+                for key, document in documents.items()
+            }
+        )
+        return await self._request(
+            "POST",
+            endpoint,
+            payload=body.model_dump(),
+            output=SummarizedModel,
+            timeout=timeout,
+        )
+
+    async def rephrase(
+        self,
+        question: str,
+        user_context: Optional[list[str]] = None,
+        context: Optional[list[Union[dict, ContextItem]]] = None,
+        model: Optional[str] = None,
+        prompt: Optional[str] = None,
+    ) -> RephraseModel:
+        endpoint = f"{self.url}{REPHRASE_PREDICT}"
+        if model:
+            endpoint += f"?model={model}"
+
+        body: dict[str, Any] = {
+            "question": question,
+            "user_context": user_context,
+            "user_id": "USER",
+        }
+        if prompt:
+            body["prompt"] = prompt
+        if context:
+            body["context"] = [
+                c.model_dump(mode="json") if isinstance(c, BaseModel) else c
+                for c in context
+            ]
+        return await self._request(
+            "POST",
+            endpoint,
+            payload=body,
+            output=RephraseModel,
+        )
+
+    async def remi(self, request: RemiRequest) -> RemiResponse:
+        endpoint = f"{self.url}{REMI_PREDICT}"
+        return await self._request(
+            "POST",
+            endpoint,
+            payload=request.model_dump(),
+            output=RemiResponse,
+        )
+
+    async def generate_retrieval(
+        self,
+        question: str,
+        context: list[str],
+        model: Optional[str] = None,
+    ) -> ChatResponse:
+        endpoint = f"{self.url}{CHAT_PREDICT}"
+        if model:
+            endpoint += f"?model={model}"
+        body = ChatModel(
+            question=question,
+            retrieval=True,
+            user_id="Nuclia PY CLI",
+            query_context=context,
+        )
+        return await self._request(
+            "POST", endpoint, payload=body.model_dump(), output=ChatResponse
+        )
+
+    async def rerank(self, model: RerankModel) -> RerankResponse:
+        endpoint = f"{self.url}{RERANK}"
+        return await self._request(
+            "POST", endpoint, payload=model.model_dump(), output=RerankResponse
+        )

hyperforge/openapi.py

@@ -0,0 +1,63 @@
+import argparse
+import datetime
+import json
+
+from fastapi import APIRouter, FastAPI
+from fastapi.openapi.utils import get_openapi
+from starlette.routing import compile_path
+
+extract_openapi_parser = argparse.ArgumentParser()
+extract_openapi_parser.add_argument("openapi_json_path", type=str)
+extract_openapi_parser.add_argument("api_version")
+extract_openapi_parser.add_argument("commit_id", type=str)
+
+
+def extract_openapi_command(component_id: str, title: str, router: APIRouter):
+    """
+    This function assumes that json, api version and commit id are coming from the command line.
+
+    However, what can be wired in here is the title of the API and the API Router
+    that provides the endpoints.
+
+    This function assumes that you want to extract things from a router in the form of:
+    - /api/v1 or /api/v2
+    """
+    args = extract_openapi_parser.parse_args()
+    openapi_json_path = args.openapi_json_path
+    api_version = args.api_version
+    commit_id = args.commit_id
+
+    app = FastAPI(title=title, version=f"{api_version}.0.0")
+
+    route_prefix = f"/api/v{api_version}"
+    routes = []
+    for route in router.routes:
+        # check if route starts with prefix and strip it, then add to new router
+        if route.path.startswith(route_prefix):  # type: ignore
+            route.path = route.path[len(route_prefix) :]  # type: ignore
+            route.path_regex, route.path_format, route.param_convertors = compile_path(  # type: ignore
+                route.path  # type: ignore
+            )
+            routes.append(route)
+
+    document = get_openapi(
+        title=app.title,
+        version=app.version,
+        openapi_version=app.openapi_version,
+        description=app.description,
+        terms_of_service=app.terms_of_service,
+        contact=app.contact,
+        license_info=app.license_info,
+        routes=routes,
+        tags=app.openapi_tags,
+        servers=app.servers,
+    )
+
+    document["x-metadata"] = {
+        component_id: {
+            "commit": commit_id,
+            "last_updated": datetime.datetime.utcnow().isoformat(),
+        }
+    }
+
+    json.dump(document, open(openapi_json_path, "w"))

hyperforge/prompts.py

@@ -0,0 +1,188 @@
+from pydantic import BaseModel
+
+from hyperforge import PROMPT_ENVIRONMENT
+
+
+class PromptArgument(BaseModel):
+    """An argument for a prompt template."""
+
+    name: str
+    """The name of the argument."""
+    description: str | None = None
+    """A human-readable description of the argument."""
+    required: bool | None = None
+    """Whether this argument must be provided."""
+
+
+class PromptConfig(BaseModel):
+    name: str
+    description: str
+    prompt: str
+    arguments: list[PromptArgument] | None = None
+    icons: dict[str, str] | None = None
+    meta: dict[str, str] | None = None
+    prompt_id: str | None = None
+
+
+VALIDATE_OR_ANSWER_PROMPT = """
+Based on the provided context and user question, perform the following tasks:
+
+1. Select only information directly relevant to the question.
+2. Break down compound sentences into simple, single-idea statements. Preserve original phrasing when possible.
+3. For any named entity with descriptive details, separate those details into distinct propositions.
+4. Ensure clarity by replacing pronouns (e.g., "it", "he", "she", "they", "this", "that") with the full names of the entities they reference, and add necessary modifiers to clarify meaning.
+5. The context may be delimited by tags such as <START OF CONTEXT> and <END OF CONTEXT>. Treat everything between these tags as context.
+6. Assess whether the context sufficiently answers the question. If it answers it partially, provide the answer; if it does not answer it fully, specify what information is missing to answer the question.
+7. If the context does not answer the question at all, just return the original question as the missing information.
+8. The `citations` field consists ONLY in a list of block IDs that are relevant to the answer, following these rules:
+  - Use the format: block-AB
+  - Just mention the block IDs, do NOT include any other text.
+  - Just mention the blocks actually relevant and that contain information used in the answer, do NOT include blocks that are not relevant.
+  - No duplicates.
+9. Do NOT hallucinate block IDs. Only use those provided in the context.
+10. Your output must be a JSON object with the following fields:
+    - "reason": Explain your reasoning for the answer or validation.
+    - "answer": Provide a partial or complete answer to the user query strictly from the information in the context. If there isn't enough information to even provide a partial answer, leave 'answer' empty.
+    - "missing_info_query": If the context is insufficient, specify what information is missing in a query shape; otherwise, leave it empty. Just return the query needed to retrieve the missing information.
+    - "useful": Indicate if the context is useful to answer the question ("yes" or "no").
+    - "citations": List the IDs of the blocks relevant to the answer, if any (e.g., ["block-AB", "block-CD"]).
+11. **IMPORTANT** If any extra instructions are provided, you MUST follow them carefully when generating the answer field. These instructions may specify the format, style, tools to use, or other requirements for the answer.
+
+{% if extra_prompts -%}
+<ADDITIONAL INSTRUCTIONS>
+These are extra instructions about how to generate the answer. They may include information about tools, style, or specific requirements.
+**IMPORTANT** You MUST use these instructions when generating the answer, and follow any specific requirements they contain.
+{% for prompt in extra_prompts -%}
+- {{ prompt }}
+{% endfor %}
+<END OF ADDITIONAL INSTRUCTIONS>
+{% endif -%}
+
+<QUESTION>
+{{question}}
+
+<START OF CONTEXT>
+{% if contexts %}
+Context:
+{% for ident, chunk in contexts.items() %}
+**{{ident}}**\n\n{{chunk}}\n\n---"
+
+{% endfor %}
+{% endif %}
+
+<END OF CONTEXT>
+
+
+"""
+
+VALIDATE_OR_ANSWER_PROMPT_TEMPLATE = PROMPT_ENVIRONMENT.from_string(
+    VALIDATE_OR_ANSWER_PROMPT
+)
+
+
+VALIDATE_JSON_SCHEMA = {
+    "title": "validate_or_answer",
+    "description": "Validate or answer",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "reason": {
+                "type": "string",
+                "description": "Reasoning for the answer or validation",
+            },
+            "answer": {
+                "type": "string",
+                "description": "Partial or complete answer to the user query from the information in the context.",
+            },
+            "missing_info_query": {
+                "type": "string",
+                "description": "Query needed to retrieve the missing information in case the context is not enough to answer the question. If the context does not answer the question at all, just return the original question.",
+            },
+            "useful": {
+                "type": "string",
+                "description": "Is the context useful to answer the question?",
+                "enum": ["yes", "no"],
+            },
+            "citations": {
+                "type": "array",
+                "items": {
+                    "type": "string",
+                    "description": "Block ID cited in the answer, e.g. block-AB",
+                },
+                "description": "List of block IDs cited in the answer, if any",
+            },
+        },
+        "required": ["reason", "answer", "missing_info_query", "useful", "citations"],
+    },
+}
+
+
+NEXT_REPHRASE_JSON_SCHEMA = {
+    "title": "next_rephrase",
+    "description": "Rephrase the question if needed given info about previous contexts and the current agent",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "rephrased_question": {
+                "type": "string",
+                "description": "Rephrased question if needed. If the question does not need to be rephrased, return an empty string.",
+            },
+            "needed": {
+                "type": "boolean",
+                "description": "Is rephrasing needed?",
+            },
+            "reason": {
+                "type": "string",
+                "description": "Reasoning for rephrasing or not the question",
+            },
+        },
+        "required": ["rephrased_question", "needed", "reason"],
+    },
+}
+
+NEXT_REPHRASE_PROMPT_SYSTEM = """
+You  decide if a question needs to be rephrased to be answered by the current agent, given the context provided by previous agents and the description of the current agent."""
+
+NEXT_REPHRASE_PROMPT = """
+You  decide if a question needs to be rephrased to be answered by the current agent, given the context provided by previous agents and the description of the current agent.
+
+Guidelines:
+1. If the question can be answered by the current agent without rephrasing, return an empty string as the rephrased question and "False" as needed.
+2. If the question needs to be rephrased to be answered by the current agent, return the rephrased question and "True" as needed.
+3. Provide a clear reasoning for your decision in the reason field.
+
+Return only a JSON object with the following fields:
+- "rephrased_question": The rephrased question if needed. If the question does not need to be rephrased, return an empty string.
+- "needed": "True" if rephrasing is needed, "False" otherwise.
+- "reason": Explain your reasoning for rephrasing or not the question.
+
+Here you have the question, the contexts from previous agents and the description of the current agent:
+
+
+<CURRENT AGENT DESCRIPTION>
+{{info}}
+
+{%if extra_info %}
+<ADDITIONAL INSTRUCTIONS>
+These are extra instructions about how to rephrase the question. They may include information about format, style, tools to use, or specific requirements.
+**IMPORTANT** You MUST use these instructions when rephrasing the question, and follow any specific requirements they contain.
+{{extra_info}}
+<END OF ADDITIONAL INSTRUCTIONS>
+{% endif %}
+
+
+<PREVIOUS CONTEXTS>
+{% if contexts %}
+Context:
+{% for context in contexts %}
+- {{context}}
+{% endfor %}
+{% else %}
+No previous context.
+{% endif %}
+
+<QUESTION>
+{{question}}
+
+"""
+NEXT_REPHRASE_PROMPT_TEMPLATE = PROMPT_ENVIRONMENT.from_string(NEXT_REPHRASE_PROMPT)

hyperforge/pubsub.py

@@ -0,0 +1,90 @@
+from typing import Annotated, Dict, Literal
+
+from pydantic import AliasChoices, BaseModel, Field
+from pydantic.types import Discriminator, Tag
+
+from hyperforge.interaction import (
+    AragAnswer,
+    Feedback,
+    OAuthAuthenticateURL,
+)
+
+# Messages used in the pubsub protocol between API and agent servers
+
+
+class StartInteraction(BaseModel):
+    """Starts an interaction for a given agent/session"""
+
+    account: str
+    agent_id: str = Field(validation_alias=AliasChoices("agent_id", "kbid"))
+    session: str
+    question_id: str
+    question: str
+    headers: Dict[str, str] = {}
+    arguments: Dict[str, str] = {}
+    workflow_id: str = "default"
+    streaming: bool = False
+    op: Literal["start"] = "start"
+
+
+class UserToAgentInteraction(BaseModel):
+    """Sends an user response"""
+
+    op: Literal["user_response"] = "user_response"
+    request_id: str
+    response: str
+
+
+class AgentPing(BaseModel):
+    """Periodic ping to indicate the agent is still working"""
+
+    op: Literal["ping"] = "ping"
+
+
+class AgentDone(BaseModel):
+    """Indicates the agent has finished and no more messages will be sent"""
+
+    op: Literal["done"] = "done"
+
+
+class AgentToUserRequest(BaseModel):
+    """Sends a partial or final answer"""
+
+    op: Literal["agent_request"] = "agent_request"
+    feedback: Feedback
+
+
+class OAuthRequest(BaseModel):
+    """Sends a partial or final answer"""
+
+    op: Literal["oauth"] = "oauth"
+    oauth: OAuthAuthenticateURL
+
+
+class AgentAnswer(BaseModel):
+    """Sends a partial or final answer"""
+
+    op: Literal["answer"] = "answer"
+    answer: AragAnswer
+
+
+# Messages send from the agent server to the API server, to be passed to the client
+AgentMessage = Annotated[
+    Annotated[AgentPing, Tag("ping")]
+    | Annotated[AgentDone, Tag("done")]
+    | Annotated[OAuthRequest, Tag("oauth")]
+    | Annotated[AgentAnswer, Tag("answer")]
+    | Annotated[AgentToUserRequest, Tag("agent_request")]
+    | Annotated[UserToAgentInteraction, Tag("user_response")],
+    Discriminator("op"),
+]
+
+
+class QuitRequest(BaseModel):
+    """Requests the agent to stop (the client has abandoned the session)"""
+
+    op: Literal["quit"] = "quit"
+
+
+# Messages sent from the API to the agent server
+APIMessage = Annotated[Annotated[QuitRequest, Tag("quit)")], Discriminator("op")]