blaxel 0.2.0rc7__tar.gz → 0.2.1rc72__tar.gz

{blaxel-0.2.0rc7 → blaxel-0.2.1rc72}/.github/workflows/test.yaml

@@ -6,6 +6,10 @@ on:
   push:
     branches: ["main", "develop"]
 
+env:
+  BL_API_KEY: fake_api_key
+  BL_WORKSPACE: fake_workspace
+
 jobs:
   test:
     name: Test blaxel package

{blaxel-0.2.0rc7 → blaxel-0.2.1rc72}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.2.0rc7
+Version: 0.2.1rc72
 Summary: Blaxel - AI development platform SDK
 Project-URL: Homepage, https://blaxel.ai
 Project-URL: Documentation, https://docs.blaxel.ai

{blaxel-0.2.0rc7 → blaxel-0.2.1rc72}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blaxel"
-version = "0.2.0rc7"
+version = "0.2.1.pre-72"
 description = "Blaxel - AI development platform SDK"
 readme = "README.md"
 authors = [{ name = "cploujoux", email = "cploujoux@blaxel.ai" }]

{blaxel-0.2.0rc7 → blaxel-0.2.1rc72}/src/blaxel/core/sandbox/sandbox.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 import time
+import uuid
 from typing import Any, Dict, List, Union
 
 from ..client.api.compute.create_sandbox import asyncio as create_sandbox
@@ -8,7 +9,7 @@ from ..client.api.compute.delete_sandbox import asyncio as delete_sandbox
 from ..client.api.compute.get_sandbox import asyncio as get_sandbox
 from ..client.api.compute.list_sandboxes import asyncio as list_sandboxes
 from ..client.client import client
-from ..client.models import Metadata, Sandbox
+from ..client.models import Metadata, Runtime, Sandbox, SandboxSpec
 from .filesystem import SandboxFileSystem
 from .network import SandboxNetwork
 from .preview import SandboxPreviews
@@ -67,14 +68,22 @@ class SandboxInstance:
                 raise Exception("Sandbox did not deploy in time")
 
     @classmethod
-    async def create(cls, sandbox: Union[Sandbox, Dict[str, Any]]) -> "SandboxInstance":
+    async def create(
+        cls, sandbox: Union[Sandbox, Dict[str, Any], None] = None
+    ) -> "SandboxInstance":
+        if sandbox is None:
+            sandbox = Sandbox()
         if isinstance(sandbox, dict):
             sandbox = Sandbox.from_dict(sandbox)
-
+        if not sandbox.metadata:
+            sandbox.metadata = Metadata(name=uuid.uuid4().hex)
         if not sandbox.spec:
-            raise Exception("Sandbox spec is required")
+            sandbox.spec = SandboxSpec(runtime=Runtime(image="blaxel/prod-base:latest"))
         if not sandbox.spec.runtime:
-            raise Exception("Sandbox runtime is required")
+            sandbox.spec.runtime = Runtime(image="blaxel/prod-base:latest")
+
+        sandbox.spec.runtime.image = sandbox.spec.runtime.image or "blaxel/prod-base:latest"
+        sandbox.spec.runtime.memory = sandbox.spec.runtime.memory or 4096
         sandbox.spec.runtime.generation = sandbox.spec.runtime.generation or "mk3"
 
         response = await create_sandbox(

{blaxel-0.2.0rc7 → blaxel-0.2.1rc72}/src/blaxel/llamaindex/custom/cohere.py

@@ -47,7 +47,6 @@ class Cohere(FunctionCallingLLM):
 
         llm = Cohere(model="command", api_key=api_key)
         resp = llm.complete("Paul Graham is ")
-        print(resp)
         ```
     """
 
@@ -55,9 +54,7 @@ class Cohere(FunctionCallingLLM):
     temperature: Optional[float] = Field(
         description="The temperature to use for sampling.", default=None
     )
-    max_retries: int = Field(
-        default=10, description="The maximum number of API retries."
-    )
+    max_retries: int = Field(default=10, description="The maximum number of API retries.")
     additional_kwargs: Dict[str, Any] = Field(
         default_factory=dict, description="Additional kwargs for the Cohere API."
     )
@@ -167,9 +164,7 @@ class Cohere(FunctionCallingLLM):
         error_on_no_tool_call: bool = False,
     ) -> List[ToolSelection]:
         """Predict and call the tool."""
-        tool_calls: List[ToolCall] = (
-            response.message.additional_kwargs.get("tool_calls", []) or []
-        )
+        tool_calls: List[ToolCall] = response.message.additional_kwargs.get("tool_calls", []) or []
 
         if len(tool_calls) < 1 and error_on_no_tool_call:
             raise ValueError(
@@ -218,10 +213,8 @@ class Cohere(FunctionCallingLLM):
 
         messages, documents = remove_documents_from_messages(messages)
 
-        tool_results: Optional[
-            List[Dict[str, Any]]
-        ] = _messages_to_cohere_tool_results_curr_chat_turn(messages) or kwargs.get(
-            "tool_results"
+        tool_results: Optional[List[Dict[str, Any]]] = (
+            _messages_to_cohere_tool_results_curr_chat_turn(messages) or kwargs.get("tool_results")
         )
         if not tool_results:
             tool_results = None
@@ -235,12 +228,8 @@ class Cohere(FunctionCallingLLM):
                 if message.role == MessageRole.TOOL:
                     temp_tool_results += _message_to_cohere_tool_results(messages, i)
 
-                    if (i == len(messages) - 1) or messages[
-                        i + 1
-                    ].role != MessageRole.TOOL:
-                        cohere_message = _get_message_cohere_format(
-                            message, temp_tool_results
-                        )
+                    if (i == len(messages) - 1) or messages[i + 1].role != MessageRole.TOOL:
+                        cohere_message = _get_message_cohere_format(message, temp_tool_results)
                         chat_history.append(cohere_message)
                         temp_tool_results = []
                 else:
@@ -262,12 +251,8 @@ class Cohere(FunctionCallingLLM):
                 if message.role == MessageRole.TOOL:
                     temp_tool_results += _message_to_cohere_tool_results(messages, i)
 
-                    if (i == len(messages) - 1) or messages[
-                        i + 1
-                    ].role != MessageRole.TOOL:
-                        cohere_message = _get_message_cohere_format(
-                            message, temp_tool_results
-                        )
+                    if (i == len(messages) - 1) or messages[i + 1].role != MessageRole.TOOL:
+                        cohere_message = _get_message_cohere_format(message, temp_tool_results)
                         chat_history.append(cohere_message)
                         temp_tool_results = []
                 else:
@@ -327,9 +312,7 @@ class Cohere(FunctionCallingLLM):
         )
 
     @llm_completion_callback()
-    def complete(
-        self, prompt: str, formatted: bool = False, **kwargs: Any
-    ) -> CompletionResponse:
+    def complete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:
         all_kwargs = self._get_all_kwargs(**kwargs)
         if "stream" in all_kwargs:
             warnings.warn(
@@ -351,9 +334,7 @@ class Cohere(FunctionCallingLLM):
         )
 
     @llm_chat_callback()
-    def stream_chat(
-        self, messages: Sequence[ChatMessage], **kwargs: Any
-    ) -> ChatResponseGen:
+    def stream_chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponseGen:
         all_kwargs = self._get_all_kwargs(**kwargs)
         all_kwargs["stream"] = True
         if all_kwargs["model"] not in CHAT_MODELS:
@@ -402,16 +383,12 @@ class Cohere(FunctionCallingLLM):
             for r in response:
                 content_delta = r.text
                 content += content_delta
-                yield CompletionResponse(
-                    text=content, delta=content_delta, raw=r._asdict()
-                )
+                yield CompletionResponse(text=content, delta=content_delta, raw=r._asdict())
 
         return gen()
 
     @llm_chat_callback()
-    async def achat(
-        self, messages: Sequence[ChatMessage], **kwargs: Any
-    ) -> ChatResponse:
+    async def achat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         all_kwargs = self._get_all_kwargs(**kwargs)
         if all_kwargs["model"] not in CHAT_MODELS:
             raise ValueError(f"{all_kwargs['model']} not supported for chat")
@@ -535,8 +512,6 @@ class Cohere(FunctionCallingLLM):
             async for r in response:
                 content_delta = r.text
                 content += content_delta
-                yield CompletionResponse(
-                    text=content, delta=content_delta, raw=r._asdict()
-                )
+                yield CompletionResponse(text=content, delta=content_delta, raw=r._asdict())
 
         return gen()

blaxel-0.2.1rc72/tests/integration/api/api_google_adk.py

@@ -0,0 +1,142 @@
+# ruff: noqa: E402
+
+import dotenv
+
+dotenv.load_dotenv()
+
+# Disable litellm logger
+import traceback
+from contextlib import asynccontextmanager
+from logging import getLogger
+from time import time
+
+import litellm
+import uvicorn
+from asgi_correlation_id import CorrelationIdMiddleware
+from fastapi import FastAPI, Request, Response, status
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse
+from google.adk.agents import Agent
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+from google.genai import types
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+from blaxel.googleadk import bl_model, bl_tools
+
+litellm._logging._disable_debugging()
+
+logger = getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logger.info("Server running on port 1338")
+    yield
+    logger.info("Server shutting down")
+
+
+MODEL = "gpt-4o-mini"
+# MODEL = "claude-3-7-sonnet-20250219"
+# MODEL = "xai-grok-beta"
+# MODEL = "cohere-command-r-plus" # x -> tool call not supported
+# MODEL = "gemini-2-5-pro-preview-03-25"
+# MODEL = "deepseek-chat"
+# MODEL = "mistral-large-latest"
+
+app = FastAPI(lifespan=lifespan)
+app.add_middleware(CorrelationIdMiddleware)
+
+
+@app.middleware("http")
+async def log_requests(request: Request, call_next):
+    start_time = time()
+
+    response: Response = await call_next(request)
+
+    process_time = (time() - start_time) * 1000
+    formatted_process_time = f"{process_time:.2f}"
+    rid_header = response.headers.get("X-Request-Id")
+    request_id = rid_header or response.headers.get("X-Blaxel-Request-Id")
+    logger.info(
+        f"{request.method} {request.url.path} {response.status_code} {formatted_process_time}ms rid={request_id}"
+    )
+
+    return response
+
+
+@app.exception_handler(Exception)
+async def validation_exception_handler(request: Request, e: Exception):
+    logger.error(f"Error on request {request.method} {request.url.path}: {e}")
+
+    # Get the full traceback information
+    tb_str = traceback.format_exception(type(e), e, e.__traceback__)
+    logger.error(f"Stacktrace: {''.join(tb_str)}")
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content=jsonable_encoder({"detail": str(e)}),
+    )
+
+
+async def weather(city: str) -> str:
+    """Get the weather in a given city"""
+    return f"The weather in {city} is sunny"
+
+
+@app.post("/")
+async def handle_request(request: Request):
+    APP_NAME = "weather_agent"
+    USER_ID = "user_1"
+    SESSION_ID = "session_001"  # Using a fixed ID for simplicity
+
+    tools = await bl_tools(["blaxel-search"], timeout_enabled=False)
+    model = await bl_model(MODEL, temperature=0)
+
+    description = "Provides weather information for specific cities."
+    prompt = """You are a helpful weather assistant. Your primary goal is to provide current weather reports. "
+When the user asks for the weather in a specific city,
+Analyze the tool's response: if the status is 'error', inform the user politely about the error message.
+If the status is 'success', present the weather 'report' clearly and concisely to the user.
+Only use the tool when a city is mentioned for a weather request.
+"""
+    agent = Agent(
+        model=model, name=APP_NAME, description=description, instruction=prompt, tools=tools
+    )
+
+    # Create the specific session where the conversation will happen
+    session_service = InMemorySessionService()
+    await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)
+    logger.info(f"Session created: App='{APP_NAME}', User='{USER_ID}', Session='{SESSION_ID}'")
+
+    runner = Runner(
+        agent=agent,
+        app_name=APP_NAME,
+        session_service=session_service,
+    )
+    logger.info(f"Runner created for agent '{runner.agent.name}'.")
+
+    input = "What is the weather like in London?"
+    content = types.Content(role="user", parts=[types.Part(text=input)])
+    final_response_text = ""
+    async for event in runner.run_async(
+        user_id=USER_ID, session_id=SESSION_ID, new_message=content
+    ):
+        # Key Concept: is_final_response() marks the concluding message for the turn.
+        if event.is_final_response():
+            if event.content and event.content.parts:
+                # Assuming text response in the first part
+                final_response_text = event.content.parts[0].text
+            elif event.actions and event.actions.escalate:  # Handle potential errors/escalations
+                final_response_text = (
+                    f"Agent escalated: {event.error_message or 'No specific message.'}"
+                )
+
+            # Add more checks here if needed (e.g., specific error codes)
+            break  # Stop processing events once the final response is found
+    return Response(final_response_text, media_type="text/plain")
+
+
+FastAPIInstrumentor.instrument_app(app, exclude_spans=["receive", "send"])
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=1338, log_level="critical", loop="asyncio")