agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/shared/ase.py

@@ -0,0 +1,54 @@
+import base64
+import hashlib
+import os
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+
+class AESCipher:
+    def __init__(self, key_string):
+        self.key = hashlib.sha256(key_string.encode()).digest()
+        self.backend = default_backend()
+
+    def encrypt(self, plaintext):
+        iv = os.urandom(16)
+        plaintext_bytes = plaintext.encode("utf-8")
+
+        encryptor = Cipher(algorithms.AES(self.key), modes.CBC(iv), backend=self.backend).encryptor()
+        padder = self._create_padder()
+
+        padded_plaintext = padder.update(plaintext_bytes) + padder.finalize()
+        ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize()
+        encrypted_data = base64.b64encode(iv + ciphertext).decode("utf-8")
+
+        return encrypted_data
+
+    def decrypt(self, encrypted_data):
+        try:
+            encrypted_bytes = base64.b64decode(encrypted_data)
+
+            iv = encrypted_bytes[:16]
+            ciphertext = encrypted_bytes[16:]
+
+            decryptor = Cipher(algorithms.AES(self.key), modes.CBC(iv), backend=self.backend).decryptor()
+            padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()
+
+            unpadder = self._create_unpadder()
+            plaintext_bytes = unpadder.update(padded_plaintext) + unpadder.finalize()
+
+            plaintext = plaintext_bytes.decode("utf-8")
+
+            return plaintext
+        except Exception as e:
+            raise ValueError(f"Decryption failed: {str(e)}")
+
+    @staticmethod
+    def _create_padder():
+        from cryptography.hazmat.primitives.padding import PKCS7
+        return PKCS7(algorithms.AES.block_size).padder()
+
+    @staticmethod
+    def _create_unpadder():
+        from cryptography.hazmat.primitives.padding import PKCS7
+        return PKCS7(algorithms.AES.block_size).unpadder()

agent_api_server/shared/base_model.py

@@ -0,0 +1,103 @@
+import logging
+from typing import Optional
+from pydantic import Field
+from pydantic import BaseModel
+from fastapi import HTTPException, status
+from typing import Any, Literal, NotRequired
+from typing_extensions import TypedDict
+
+logger = logging.getLogger(__name__)
+
+class ThreadInfo(BaseModel):
+    thread_id: str
+    graph_name: str
+    status: str
+
+class RunResponse(BaseModel):
+    status: str
+    thread_id: str
+    result: Optional[Any] = None
+    error: Optional[str] = None
+
+def error_response(
+    status_code: int,
+    error_type: str,
+    message: str,
+    **additional_info: Any
+) -> HTTPException:
+    """Helper function to create consistent error responses."""
+    detail = {
+        "status": "error",
+        "error": error_type,
+        "message": message,
+        **additional_info
+    }
+    return HTTPException(
+        status_code=status_code,
+        detail=detail
+    )
+
+class ToolCall(TypedDict):
+    """Represents a request to call a tool."""
+
+    name: str
+    """The name of the tool to be called."""
+    args: dict[str, Any]
+    """The arguments to the tool call."""
+    id: str | None
+    """An identifier associated with the tool call."""
+    type: NotRequired[Literal["tool_call"]]
+
+class ChatMessage(BaseModel):
+    """Message in a chat."""
+    type: Literal["human", "ai", "tool"] = Field(
+        description="Role of the message.",
+        examples=["human", "ai", "tool"],
+    )
+    content_type: Literal["markdown", "json", "html", "xml", "python", "yaml", "text"] = Field(
+        description="The data type of the message.",
+        examples=["markdown", "json", "html", "xml", "python", "yaml", "text"],
+    )
+    content: str = Field(
+        description="Content of the message.",
+        examples=["Hello, world!"],
+    )
+    tool_calls: list[ToolCall] = Field(
+        description="Tool calls in the message.",
+        default=[],
+    )
+    response_metadata: dict[str, Any] = Field(
+        description="Response metadata. For example: response headers, logprobs, token counts.",
+        default={},
+    )
+    references: list[dict[str, Any]] = Field(
+        description="references metadata. For example: List of reference materials.",
+        default=[],
+    )
+
+    def pretty_repr(self) -> str:
+        """Get a pretty representation of the message."""
+        base_title = self.type.title() + " Message"
+        padded = " " + base_title + " "
+        sep_len = (80 - len(padded)) // 2
+        sep = "=" * sep_len
+        second_sep = sep + "=" if len(padded) % 2 else sep
+        title = f"{sep}{padded}{second_sep}"
+        return f"{title}\n\n{self.content}"
+
+    def pretty_print(self) -> None:
+        logger.info(self.pretty_repr())  # noqa: T201
+
+
+def sse_response_example() -> dict[int | str, Any]:
+    return {
+        status.HTTP_200_OK: {
+            "description": "Server Sent Event Response",
+            "content": {
+                "text/event-stream": {
+                    "example": "data: {'type': 'token', 'content': 'Hello'}\n\ndata: {'type': 'token', 'content': ' World'}\n\ndata: [DONE]\n\n",
+                    "schema": {"type": "string"},
+                }
+            },
+        }
+    }

agent_api_server/shared/common.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+from typing import Dict, Set
+from enum import Enum
+
+
+class ConfigCategory(Enum):
+    CHAT_PROVIDER = "CHAT_PROVIDER"
+    CHAT_MODEL = "CHAT_MODEL"
+    CHAT_CREDENTIALS = "CHAT_CREDENTIALS"
+
+    EMBEDDING_PROVIDER = "EMBEDDING_PROVIDER"
+    EMBEDDING_MODEL = "EMBEDDING_MODEL"
+    EMBEDDING_CREDENTIALS = "EMBEDDING_CREDENTIALS"
+
+    RERANK_PROVIDER = "RERANK_PROVIDER"
+    RERANK_MODEL = "RERANK_MODEL"
+    RERANK_CREDENTIALS = "RERANK_CREDENTIALS"
+
+    AGENT_ID = "AGENT_ID"
+
+MODEL_SUFFIX_MAP = {
+    "CHAT_PROVIDER": "llm",
+    "EMBEDDING_PROVIDER": "text-embedding",
+    "RERANK_PROVIDER": "rerank",
+}
+
+PROVIDER_FIELDS = {
+    ConfigCategory.CHAT_PROVIDER.value: "llm",
+    ConfigCategory.EMBEDDING_PROVIDER.value: "text-embedding",
+    ConfigCategory.RERANK_PROVIDER.value: "rerank"
+}
+
+def process_model_from_config_dict(input_dict: Dict) -> Dict:
+    result = {
+        "models": {},
+        "required": {}
+    }
+
+    if input_dict is None:
+        return result
+
+    properties = input_dict.get('properties', {})
+
+    tool_model_map: Dict[str, Set[str]] = {}
+
+    for prop_name, prop_def in properties.items():
+        prop_name_upper = prop_name.upper()
+
+        if prop_name_upper in PROVIDER_FIELDS:
+            tool_name = "default"
+            model_type = PROVIDER_FIELDS[prop_name_upper]
+            if tool_name not in tool_model_map:
+                tool_model_map[tool_name] = set()
+            tool_model_map[tool_name].add(model_type)
+            continue
+
+        if prop_name_upper.endswith("_CHAT_PROVIDER"):
+            tool_name = prop_name_upper[:-len("_CHAT_PROVIDER")].lower()
+            if not tool_name:
+                tool_name = "default"
+            model_type = "llm"
+        elif prop_name_upper.endswith("_EMBEDDING_PROVIDER"):
+            tool_name = prop_name_upper[:-len("_EMBEDDING_PROVIDER")].lower()
+            if not tool_name:
+                tool_name = "default"
+            model_type = "text-embedding"
+        elif prop_name_upper.endswith("_RERANK_PROVIDER"):
+            tool_name = prop_name_upper[:-len("_RERANK_PROVIDER")].lower()
+            if not tool_name:
+                tool_name = "default"
+            model_type = "rerank"
+        else:
+            continue
+
+        if tool_name:
+            if tool_name not in tool_model_map:
+                tool_model_map[tool_name] = set()
+            tool_model_map[tool_name].add(model_type)
+
+    standard_fields = ['CHAT_PROVIDER', 'EMBEDDING_PROVIDER']
+    for field in standard_fields:
+        if field in properties:
+            tool_name = "default"
+            model_type = "llm" if field == "CHAT_PROVIDER" else "text-embedding"
+            if tool_name not in tool_model_map:
+                tool_model_map[tool_name] = set()
+            tool_model_map[tool_name].add(model_type)
+
+    if not tool_model_map:
+        for prop_name in properties.keys():
+            prop_name_upper = prop_name.upper()
+            if 'CHAT' in prop_name_upper and 'PROVIDER' in prop_name_upper:
+                tool_name = "default"
+                model_type = "llm"
+                if tool_name not in tool_model_map:
+                    tool_model_map[tool_name] = set()
+                tool_model_map[tool_name].add(model_type)
+            elif 'EMBEDDING' in prop_name_upper and 'PROVIDER' in prop_name_upper:
+                tool_name = "default"
+                model_type = "text-embedding"
+                if tool_name not in tool_model_map:
+                    tool_model_map[tool_name] = set()
+                tool_model_map[tool_name].add(model_type)
+
+    for tool, model_types in tool_model_map.items():
+        model_list = list(model_types)
+        result['models'][tool] = model_list
+        result['required'][tool] = model_list.copy()
+
+    return result

agent_api_server/shared/decode_token.py

@@ -0,0 +1,107 @@
+import base64
+import json
+import logging
+from datetime import datetime
+from typing import Dict, Any, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+def decode_jwt(token: str) -> Tuple[Optional[Dict], Optional[Dict], Optional[Dict]]:
+    try:
+        parts = token.split('.')
+
+        if len(parts) != 3:
+            logger.error("Invalid JWT format: expected 3 parts, got %d", len(parts))
+            return None, None, None
+
+        header = _decode_base64_json(parts[0], "header")
+        if not header:
+            return None, None, None
+
+        payload = _decode_base64_json(parts[1], "payload")
+        if not payload:
+            return None, None, None
+
+        user_info = _extract_user_info(payload)
+
+        _log_jwt_info(header, payload, parts[2])
+
+        return header, payload, user_info
+
+    except Exception as e:
+        logger.error("Failed to decode JWT: %s", str(e))
+        return None, None, None
+
+
+def _decode_base64_json(encoded_str: str, component_name: str) -> Optional[Dict]:
+    """Decode base64 URL safe string to JSON dictionary."""
+    try:
+        # Pad with '=' for proper base64 decoding
+        padding_needed = len(encoded_str) % 4
+        if padding_needed:
+            encoded_str += '=' * (4 - padding_needed)
+
+        decoded_bytes = base64.urlsafe_b64decode(encoded_str)
+        decoded_dict = json.loads(decoded_bytes)
+        return decoded_dict
+
+    except Exception as e:
+        logger.error("Failed to decode %s: %s", component_name, str(e))
+        return None
+
+
+def _extract_user_info(payload: Dict) -> Dict[str, Any]:
+    user_info = {}
+
+    user_id = payload.get('id')
+    if user_id:
+        user_info['id'] = user_id
+        logger.info("Extracted user ID: %s", user_id)
+    else:
+        logger.warning("User ID not found in payload")
+
+    username = payload.get('username')
+    if username:
+        user_info['username'] = username
+        logger.info("Extracted username: %s", username)
+    else:
+        logger.warning("Username not found in payload")
+
+    return user_info
+
+
+def _log_jwt_info(header: Dict, payload: Dict, signature: str) -> None:
+    logger.info("JWT Decoded Successfully")
+    logger.info("=" * 50)
+
+    logger.info("Header:")
+    logger.info(json.dumps(header, indent=2))
+    logger.info("")
+
+    logger.info("Payload:")
+    for key, value in payload.items():
+        if key in ['creationTime', 'exp', 'iat', 'lastModifiedTime'] and isinstance(value, (int, float)):
+            try:
+                dt = datetime.fromtimestamp(value)
+                logger.info("  %s: %s (%s)", key, value, dt.strftime('%Y-%m-%d %H:%M:%S'))
+            except Exception:
+                logger.info("  %s: %s", key, value)
+        else:
+            logger.info("  %s: %s", key, value)
+
+    signature_preview = signature[:50] + "..." if len(signature) > 50 else signature
+    logger.info("\nSignature: %s", signature_preview)
+    logger.info("=" * 50)
+
+
+if __name__ == "__main__":
+    # Example token (truncated for demonstration)
+    example_token = "eyJhbGciOiJIUzUxMiIsInR5cCI6IkpXVCJ9.eyJjb3VudHJ5IjoiQ04iLCJjcmVhdGlvblRpbWUiOjE3NjQxNDgwNTMsImV4cCI6MTc3MjI2MDYxNywiZmlyc3ROYW1lIjoiYWRtaW4iLCJpYXQiOjE3NzIyNTcwMTcsImlkIjoiZTQ2MDk2YWUtODhmMC00NDIxLWI1YmUtMTU3OTAxZmRhZDZkIiwiaXNzIjoid2lzZS1wYWFzIiwibGFzdE1vZGlmaWVkVGltZSI6MTc3MjE2MDg3OCwibGFzdE5hbWUiOiJhZG1pbiIsImxvZ2luIjoiYWRtaW5AaW90LnNlbnNlIiwicmVmcmVzaFRva2VuIjoiN2Y4ZTkzNmMtMTQ2Ny0xMWYxLThiODUtZTJlZTMyNTQ2MWJlIiwic3RhdHVzIjoiQWN0aXZlIiwidXNlcm5hbWUiOiJhZG1pbkBpb3Quc2Vuc2UifQ.cmxJ4OG6xsKKK0qom2ur84qmra0P1_nGmCvEM7YRThN8ECyJNXQaLnXYhzEHYjbtr9anBuUoMwqkIc1kZg7t2Q"
+
+    header, payload, user_info = decode_jwt(example_token)
+
+    if user_info:
+        print(f"\nExtracted User Information:")
+        print(f"User ID: {user_info.get('id', 'Not found')}")
+        print(f"Username: {user_info.get('username', 'Not found')}")