From cc53872a6a065d17e8cc60fc49028449703e5f5a Mon Sep 17 00:00:00 2001 From: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Fri, 7 Mar 2025 16:15:51 +0530 Subject: [PATCH 1/2] Added headers to prompt service calls Signed-off-by: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com> --- src/unstract/sdk/prompt.py | 41 ++++++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/src/unstract/sdk/prompt.py b/src/unstract/sdk/prompt.py index 0d908e6c..c4ebb4d7 100644 --- a/src/unstract/sdk/prompt.py +++ b/src/unstract/sdk/prompt.py @@ -36,36 +36,50 @@ def __init__( @log_elapsed(operation="ANSWER_PROMPTS") def answer_prompt( - self, payload: dict[str, Any], params: Optional[dict[str, str]] = None + self, + payload: dict[str, Any], + params: Optional[dict[str, str]] = None, + headers: Optional[dict[str, str]] = None, ) -> dict[str, Any]: url_path = "answer-prompt" if self.is_public_call: url_path = "answer-prompt-public" return self._post_call( - url_path=url_path, - payload=payload, - params=params, + url_path=url_path, payload=payload, params=params, headers=headers ) def single_pass_extraction( - self, payload: dict[str, Any], params: Optional[dict[str, str]] = None + self, + payload: dict[str, Any], + params: Optional[dict[str, str]] = None, + headers: Optional[dict[str, str]] = None, ) -> dict[str, Any]: return self._post_call( url_path="single-pass-extraction", payload=payload, params=params, + headers=headers, ) def summarize( - self, payload: dict[str, Any], params: Optional[dict[str, str]] = None + self, + payload: dict[str, Any], + params: Optional[dict[str, str]] = None, + headers: Optional[dict[str, str]] = None, ) -> dict[str, Any]: - return self._post_call(url_path="summarize", payload=payload, params=params) + return self._post_call( + url_path="summarize", + payload=payload, + params=params, + headers=headers, + ) def _post_call( self, url_path: str, payload: dict[str, Any], params: Optional[dict[str, str]] = None, + headers: Optional[dict[str, str]] = None, ) -> dict[str, Any]: """Invokes and communicates to prompt service to fetch response for the prompt. @@ -74,6 +88,7 @@ def _post_call( url_path (str): URL path to the service endpoint payload (dict): Payload to send in the request body params (dict, optional): Query parameters to include in the request + headers (dict, optional): Headers to include in the request Returns: dict: Response from the prompt service @@ -94,13 +109,19 @@ def _post_call( "status_code": 500, } url: str = f"{self.base_url}/{url_path}" - headers: dict[str, str] = {} + + default_headers = {} + if not self.is_public_call: - headers = {"Authorization": f"Bearer {self.bearer_token}"} + default_headers = {"Authorization": f"Bearer {self.bearer_token}"} + + if headers: + default_headers.update(headers) + response: Response = Response() try: response = requests.post( - url=url, json=payload, params=params, headers=headers + url=url, json=payload, params=params, headers=default_headers ) response.raise_for_status() result["status"] = "OK" From c6f8436e610e9b0323558ec026554f45fa210e63 Mon Sep 17 00:00:00 2001 From: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Fri, 7 Mar 2025 16:16:15 +0530 Subject: [PATCH 2/2] Fixed pre-commit issues Signed-off-by: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com> --- src/unstract/sdk/adapters/embedding/open_ai/src/open_ai.py | 1 - .../embedding/qdrant_fast_embed/src/qdrant_fast_embed.py | 1 - src/unstract/sdk/adapters/embedding/vertex_ai/src/vertex_ai.py | 3 +-- .../adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py | 2 -- 4 files changed, 1 insertion(+), 6 deletions(-) diff --git a/src/unstract/sdk/adapters/embedding/open_ai/src/open_ai.py b/src/unstract/sdk/adapters/embedding/open_ai/src/open_ai.py index 875e4bd6..0ca1ce7f 100644 --- a/src/unstract/sdk/adapters/embedding/open_ai/src/open_ai.py +++ b/src/unstract/sdk/adapters/embedding/open_ai/src/open_ai.py @@ -6,7 +6,6 @@ from llama_index.embeddings.openai import OpenAIEmbedding from unstract.sdk.adapters.embedding.embedding_adapter import EmbeddingAdapter -from unstract.sdk.adapters.embedding.helper import EmbeddingHelper from unstract.sdk.adapters.exceptions import AdapterError diff --git a/src/unstract/sdk/adapters/embedding/qdrant_fast_embed/src/qdrant_fast_embed.py b/src/unstract/sdk/adapters/embedding/qdrant_fast_embed/src/qdrant_fast_embed.py index 2329c1c9..94bb4d29 100644 --- a/src/unstract/sdk/adapters/embedding/qdrant_fast_embed/src/qdrant_fast_embed.py +++ b/src/unstract/sdk/adapters/embedding/qdrant_fast_embed/src/qdrant_fast_embed.py @@ -5,7 +5,6 @@ from llama_index.embeddings.fastembed import FastEmbedEmbedding from unstract.sdk.adapters.embedding.embedding_adapter import EmbeddingAdapter -from unstract.sdk.adapters.embedding.helper import EmbeddingHelper from unstract.sdk.adapters.exceptions import AdapterError diff --git a/src/unstract/sdk/adapters/embedding/vertex_ai/src/vertex_ai.py b/src/unstract/sdk/adapters/embedding/vertex_ai/src/vertex_ai.py index c2287879..395329e6 100644 --- a/src/unstract/sdk/adapters/embedding/vertex_ai/src/vertex_ai.py +++ b/src/unstract/sdk/adapters/embedding/vertex_ai/src/vertex_ai.py @@ -1,7 +1,6 @@ import json -import logging import os -from typing import Any, Optional +from typing import Any from google.auth.transport import requests as google_requests from google.oauth2.service_account import Credentials diff --git a/src/unstract/sdk/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py b/src/unstract/sdk/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py index 6b06a9c6..2b5b06b1 100644 --- a/src/unstract/sdk/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py +++ b/src/unstract/sdk/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py @@ -1,4 +1,3 @@ -import json import logging import os from typing import Any, Optional @@ -11,7 +10,6 @@ TextExtractionResult, ) from unstract.sdk.adapters.x2text.llm_whisperer_v2.src.constants import ( - HTTPMethod, WhispererEndpoint, ) from unstract.sdk.adapters.x2text.llm_whisperer_v2.src.dto import WhispererRequestParams