v2.12.0 (#725)
Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Sirieix
GitHub
parent
cf3f4ca58f
commit
1843196d88
20 changed files with 399 additions and 1675 deletions
|
|
@ -14,7 +14,6 @@ import time
|
|||
from typing import Any, ClassVar, assert_never, cast
|
||||
from uuid import uuid4
|
||||
from weakref import WeakKeyDictionary
|
||||
import webbrowser
|
||||
|
||||
from pydantic import BaseModel
|
||||
from rich import print as rprint
|
||||
|
|
@ -160,16 +159,12 @@ from vibe.core.session.title_format import format_session_title
|
|||
from vibe.core.skills.manager import SkillManager
|
||||
from vibe.core.teleport.telemetry import send_teleport_early_failure_telemetry
|
||||
from vibe.core.teleport.types import (
|
||||
TeleportAuthCompleteEvent,
|
||||
TeleportAuthRequiredEvent,
|
||||
TeleportCheckingGitEvent,
|
||||
TeleportCompleteEvent,
|
||||
TeleportFetchingUrlEvent,
|
||||
TeleportPushingEvent,
|
||||
TeleportPushRequiredEvent,
|
||||
TeleportPushResponseEvent,
|
||||
TeleportStartingWorkflowEvent,
|
||||
TeleportWaitingForGitHubEvent,
|
||||
)
|
||||
from vibe.core.tools.builtins.ask_user_question import (
|
||||
AskUserQuestionArgs,
|
||||
|
|
@ -1604,15 +1599,6 @@ class VibeApp(App): # noqa: PLR0904
|
|||
teleport_msg.set_status("Syncing with remote...")
|
||||
case TeleportStartingWorkflowEvent():
|
||||
teleport_msg.set_status("Teleporting...")
|
||||
case TeleportWaitingForGitHubEvent(message=msg):
|
||||
teleport_msg.set_status(msg or "Connecting to GitHub...")
|
||||
case TeleportAuthRequiredEvent(oauth_url=url, message=msg):
|
||||
webbrowser.open(url)
|
||||
teleport_msg.set_status(msg or "Authorizing GitHub...")
|
||||
case TeleportAuthCompleteEvent():
|
||||
teleport_msg.set_status("GitHub authorized")
|
||||
case TeleportFetchingUrlEvent():
|
||||
teleport_msg.set_status("Finalizing...")
|
||||
case TeleportCompleteEvent(url=url):
|
||||
teleport_msg.set_complete(url)
|
||||
except TeleportError as e:
|
||||
|
|
|
|||
|
|
@ -675,10 +675,8 @@ class AgentLoop: # noqa: PLR0904
|
|||
raise TeleportError("_TeleportService is unexpectedly None")
|
||||
self._teleport_service = _TeleportService(
|
||||
session_logger=self.session_logger,
|
||||
vibe_code_base_url=self.config.vibe_code_base_url,
|
||||
vibe_code_workflow_id=self.config.vibe_code_workflow_id,
|
||||
vibe_code_sessions_base_url=self.config.vibe_code_sessions_base_url,
|
||||
vibe_code_api_key=self.config.vibe_code_api_key,
|
||||
vibe_code_task_queue=self.config.vibe_code_task_queue,
|
||||
vibe_config=self._base_config,
|
||||
)
|
||||
return self._teleport_service
|
||||
|
|
@ -687,29 +685,23 @@ class AgentLoop: # noqa: PLR0904
|
|||
async def teleport_to_vibe_code(
|
||||
self, prompt: str | None
|
||||
) -> AsyncGenerator[TeleportYieldEvent, TeleportPushResponseEvent | None]:
|
||||
from vibe.core.teleport.nuage import TeleportSession
|
||||
|
||||
session_messages = [
|
||||
msg.model_dump(exclude_none=True) for msg in self.messages[1:]
|
||||
]
|
||||
nb_session_messages = max(len(self.messages) - 1, 0)
|
||||
if prompt:
|
||||
resolved_prompt = prompt
|
||||
else:
|
||||
last = self._last_user_message()
|
||||
content = last.content if last else None
|
||||
resolved_prompt = (
|
||||
f"{content} (continue)" if isinstance(content, str) and content else ""
|
||||
)
|
||||
telemetry_tracker = TeleportTelemetryTracker(
|
||||
telemetry_client=self.telemetry_client,
|
||||
nb_session_messages=len(session_messages),
|
||||
stage="no_history"
|
||||
if prompt is None and not session_messages
|
||||
else "git_check",
|
||||
)
|
||||
session = TeleportSession(
|
||||
metadata={
|
||||
"agent": self.agent_profile.name,
|
||||
"model": self.config.active_model,
|
||||
"stats": self.stats.model_dump(),
|
||||
},
|
||||
messages=session_messages,
|
||||
nb_session_messages=nb_session_messages,
|
||||
stage="no_history" if not resolved_prompt else "git_check",
|
||||
)
|
||||
try:
|
||||
async with self.teleport_service:
|
||||
gen = self.teleport_service.execute(prompt=prompt, session=session)
|
||||
gen = self.teleport_service.execute(prompt=resolved_prompt)
|
||||
response: TeleportPushResponseEvent | None = None
|
||||
while True:
|
||||
try:
|
||||
|
|
@ -733,6 +725,16 @@ class AgentLoop: # noqa: PLR0904
|
|||
telemetry_tracker.send_failure_if_needed()
|
||||
self._teleport_service = None
|
||||
|
||||
def _last_user_message(self) -> LLMMessage | None:
|
||||
return next(
|
||||
(
|
||||
m
|
||||
for m in reversed(self.messages)
|
||||
if m.role == Role.user and not m.injected
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
def _setup_middleware(self) -> None:
|
||||
"""Configure middleware pipeline for this conversation."""
|
||||
self.middleware_pipeline.clear()
|
||||
|
|
@ -1330,14 +1332,7 @@ class AgentLoop: # noqa: PLR0904
|
|||
available_tools = self.format_handler.get_available_tools(self.tool_manager)
|
||||
tool_choice = self.format_handler.get_tool_choice()
|
||||
|
||||
last_user_message = next(
|
||||
(
|
||||
m
|
||||
for m in reversed(self.messages)
|
||||
if m.role == Role.user and not m.injected
|
||||
),
|
||||
None,
|
||||
)
|
||||
last_user_message = self._last_user_message()
|
||||
self.telemetry_client.send_request_sent(
|
||||
model=active_model.alias,
|
||||
nb_context_chars=sum(len(m.content or "") for m in self.messages),
|
||||
|
|
@ -1400,14 +1395,7 @@ class AgentLoop: # noqa: PLR0904
|
|||
available_tools = self.format_handler.get_available_tools(self.tool_manager)
|
||||
tool_choice = self.format_handler.get_tool_choice()
|
||||
|
||||
last_user_message = next(
|
||||
(
|
||||
m
|
||||
for m in reversed(self.messages)
|
||||
if m.role == Role.user and not m.injected
|
||||
),
|
||||
None,
|
||||
)
|
||||
last_user_message = self._last_user_message()
|
||||
self.telemetry_client.send_request_sent(
|
||||
model=active_model.alias,
|
||||
nb_context_chars=sum(len(m.content or "") for m in self.messages),
|
||||
|
|
|
|||
|
|
@ -525,11 +525,12 @@ class VibeConfig(BaseSettings):
|
|||
|
||||
vibe_code_enabled: bool = Field(default=True, exclude=True)
|
||||
vibe_code_base_url: str = Field(default="https://api.mistral.ai", exclude=True)
|
||||
vibe_code_sessions_base_url: str = Field(
|
||||
default="https://chat.mistral.ai", exclude=True
|
||||
)
|
||||
vibe_code_workflow_id: str = Field(default="__shared-nuage-workflow", exclude=True)
|
||||
vibe_code_task_queue: str | None = Field(default="shared-vibe-nuage", exclude=True)
|
||||
vibe_code_api_key_env_var: str = Field(default="MISTRAL_API_KEY", exclude=True)
|
||||
vibe_code_project_name: str | None = Field(default=None, exclude=True)
|
||||
vibe_code_experimental_nuage_enabled: bool = Field(default=False, exclude=True)
|
||||
|
||||
# TODO(otel): remove exclude=True once the feature is publicly available
|
||||
enable_otel: bool = Field(default=False, exclude=True)
|
||||
|
|
|
|||
|
|
@ -6,15 +6,11 @@ import sys
|
|||
from typing import TextIO
|
||||
|
||||
from vibe.core.teleport.types import (
|
||||
TeleportAuthCompleteEvent,
|
||||
TeleportAuthRequiredEvent,
|
||||
TeleportCheckingGitEvent,
|
||||
TeleportCompleteEvent,
|
||||
TeleportFetchingUrlEvent,
|
||||
TeleportPushingEvent,
|
||||
TeleportPushRequiredEvent,
|
||||
TeleportStartingWorkflowEvent,
|
||||
TeleportWaitingForGitHubEvent,
|
||||
)
|
||||
from vibe.core.types import AssistantEvent, BaseEvent, LLMMessage, OutputFormat
|
||||
|
||||
|
|
@ -62,14 +58,6 @@ class TextOutputFormatter(OutputFormatter):
|
|||
self._print("Syncing with remote...")
|
||||
case TeleportStartingWorkflowEvent():
|
||||
self._print("Teleporting...")
|
||||
case TeleportWaitingForGitHubEvent(message=msg):
|
||||
self._print(msg or "Connecting to GitHub...")
|
||||
case TeleportAuthRequiredEvent(oauth_url=url, message=msg):
|
||||
self._print(msg or f"Open to authorize GitHub: {url}")
|
||||
case TeleportAuthCompleteEvent():
|
||||
self._print("GitHub authorized")
|
||||
case TeleportFetchingUrlEvent():
|
||||
self._print("Finalizing...")
|
||||
case TeleportCompleteEvent():
|
||||
self._final_response = event.url
|
||||
|
||||
|
|
|
|||
|
|
@ -349,15 +349,10 @@ class TelemetryClient:
|
|||
)
|
||||
|
||||
def send_teleport_completed(
|
||||
self,
|
||||
*,
|
||||
push_required: bool,
|
||||
github_auth_required: bool,
|
||||
nb_session_messages: int,
|
||||
self, *, push_required: bool, nb_session_messages: int
|
||||
) -> None:
|
||||
payload: TeleportCompletedPayload = {
|
||||
"push_required": push_required,
|
||||
"github_auth_required": github_auth_required,
|
||||
"nb_session_messages": nb_session_messages,
|
||||
}
|
||||
self.send_telemetry_event("vibe.teleport_completed", dict(payload))
|
||||
|
|
@ -368,14 +363,12 @@ class TelemetryClient:
|
|||
stage: TeleportFailureStage,
|
||||
error_class: str,
|
||||
push_required: bool,
|
||||
github_auth_required: bool,
|
||||
nb_session_messages: int,
|
||||
) -> None:
|
||||
payload: TeleportFailedPayload = {
|
||||
"stage": stage,
|
||||
"error_class": error_class,
|
||||
"push_required": push_required,
|
||||
"github_auth_required": github_auth_required,
|
||||
"nb_session_messages": nb_session_messages,
|
||||
}
|
||||
self.send_telemetry_event("vibe.teleport_failed", dict(payload))
|
||||
|
|
|
|||
|
|
@ -40,20 +40,12 @@ class TelemetryRequestMetadata(TelemetryBaseMetadata):
|
|||
|
||||
|
||||
TeleportFailureStage = Literal[
|
||||
"no_history",
|
||||
"remote_session",
|
||||
"git_check",
|
||||
"push",
|
||||
"workflow_start",
|
||||
"github_auth",
|
||||
"fetch_url",
|
||||
"cancelled",
|
||||
"no_history", "remote_session", "git_check", "push", "workflow_start", "cancelled"
|
||||
]
|
||||
|
||||
|
||||
class TeleportCompletedPayload(TypedDict):
|
||||
push_required: bool
|
||||
github_auth_required: bool
|
||||
nb_session_messages: int
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,132 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
from typing import Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.utils.http import build_ssl_context
|
||||
|
||||
|
||||
class ExperimentalNuageTextPart(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
type: str = "text"
|
||||
text: str
|
||||
|
||||
|
||||
class ExperimentalNuageMessage(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
role: str = "user"
|
||||
parts: list[ExperimentalNuageTextPart]
|
||||
|
||||
|
||||
class ExperimentalNuageDiff(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
format: Literal["git-diff"] = "git-diff"
|
||||
encoding: Literal["base64"] = "base64"
|
||||
compression: Literal["zstd"] = "zstd"
|
||||
content: str
|
||||
|
||||
|
||||
class ExperimentalNuageRepository(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
repo_url: str = Field(serialization_alias="repoUrl")
|
||||
branch: str | None = None
|
||||
commit_sha: str | None = Field(default=None, serialization_alias="commitSha")
|
||||
diff: ExperimentalNuageDiff | None = None
|
||||
|
||||
|
||||
class ExperimentalNuageContext(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
repositories: list[ExperimentalNuageRepository]
|
||||
|
||||
|
||||
class ExperimentalNuageRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
project_name: str = Field(default="Vibe CLI", serialization_alias="project_name")
|
||||
source: str = "vibe_code_cli"
|
||||
idempotency_key: str = Field(serialization_alias="idempotencyKey")
|
||||
message: ExperimentalNuageMessage
|
||||
context: ExperimentalNuageContext
|
||||
|
||||
|
||||
class ExperimentalNuageResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
nuage_session_id: str = Field(validation_alias="sessionId")
|
||||
nuage_web_session_id: str = Field(validation_alias="webSessionId")
|
||||
nuage_project_id: str = Field(validation_alias="projectId")
|
||||
status: str
|
||||
url: str
|
||||
|
||||
|
||||
class ExperimentalNuageClient:
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
*,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._client = client
|
||||
self._owns_client = client is None
|
||||
self._timeout = timeout
|
||||
|
||||
async def __aenter__(self) -> ExperimentalNuageClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._timeout), verify=build_ssl_context()
|
||||
)
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: types.TracebackType | None,
|
||||
) -> None:
|
||||
if self._owns_client and self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
@property
|
||||
def _http_client(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._timeout), verify=build_ssl_context()
|
||||
)
|
||||
self._owns_client = True
|
||||
return self._client
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async def start(
|
||||
self, request: ExperimentalNuageRequest
|
||||
) -> ExperimentalNuageResponse:
|
||||
response = await self._http_client.post(
|
||||
f"{self._base_url}/api/v1/code/sessions",
|
||||
headers=self._headers(),
|
||||
json=request.model_dump(mode="json", by_alias=True, exclude_none=True),
|
||||
)
|
||||
if not response.is_success:
|
||||
raise ServiceTeleportError(f"Vibe Code Nuage start failed: {response.text}")
|
||||
|
||||
try:
|
||||
return ExperimentalNuageResponse.model_validate(response.json())
|
||||
except (ValueError, ValidationError) as e:
|
||||
raise ServiceTeleportError("Vibe Code Nuage response was invalid") from e
|
||||
|
|
@ -1,104 +1,71 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from enum import StrEnum, auto
|
||||
import time
|
||||
import types
|
||||
from typing import Any
|
||||
from typing import Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.utils.http import build_ssl_context
|
||||
|
||||
|
||||
class GitHubParams(BaseModel):
|
||||
repo: str | None = None
|
||||
branch: str | None = None
|
||||
commit: str | None = None
|
||||
pr_number: int | None = None
|
||||
teleported_diffs: bytes | None = None
|
||||
class NuageTextPart(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class ChatAssistantParams(BaseModel):
|
||||
create_thread: bool = False
|
||||
user_message: str | None = None
|
||||
project_name: str | None = None
|
||||
|
||||
|
||||
class TeleportSession(BaseModel):
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
messages: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
|
||||
class WorkflowIntegrations(BaseModel):
|
||||
github: GitHubParams | None = None
|
||||
chat_assistant: ChatAssistantParams | None = None
|
||||
|
||||
|
||||
class VibeAgent(BaseModel):
|
||||
polymorphic_type: str = "vibe_agent"
|
||||
name: str = "vibe-agent"
|
||||
vibe_config: dict[str, Any] | None = None
|
||||
session: TeleportSession | None = None
|
||||
|
||||
|
||||
class WorkflowConfig(BaseModel):
|
||||
agent: VibeAgent = Field(default_factory=VibeAgent)
|
||||
|
||||
|
||||
class TextChunk(BaseModel):
|
||||
type: str = "text"
|
||||
text: str
|
||||
|
||||
|
||||
class WorkflowParams(BaseModel):
|
||||
prompt: str
|
||||
message: list[TextChunk] | None = None
|
||||
config: WorkflowConfig = Field(default_factory=WorkflowConfig)
|
||||
integrations: WorkflowIntegrations = Field(default_factory=WorkflowIntegrations)
|
||||
class NuageMessage(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
role: str = "user"
|
||||
parts: list[NuageTextPart]
|
||||
|
||||
|
||||
class WorkflowExecuteResponse(BaseModel):
|
||||
execution_id: str
|
||||
class NuageDiff(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
format: Literal["git-diff"] = "git-diff"
|
||||
encoding: Literal["base64"] = "base64"
|
||||
compression: Literal["zstd"] = "zstd"
|
||||
content: str
|
||||
|
||||
|
||||
class GitHubStatus(StrEnum):
|
||||
PENDING = auto()
|
||||
WAITING_FOR_OAUTH = auto()
|
||||
CONNECTED = auto()
|
||||
OAUTH_TIMEOUT = auto()
|
||||
ERROR = auto()
|
||||
class NuageRepository(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
repo_url: str = Field(serialization_alias="repoUrl")
|
||||
branch: str | None = None
|
||||
commit_sha: str | None = Field(default=None, serialization_alias="commitSha")
|
||||
diff: NuageDiff | None = None
|
||||
|
||||
|
||||
class GitHubPublicData(BaseModel):
|
||||
status: GitHubStatus
|
||||
oauth_url: str | None = None
|
||||
error: str | None = None
|
||||
working_branch: str | None = None
|
||||
repo: str | None = None
|
||||
class NuageContext(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
@property
|
||||
def connected(self) -> bool:
|
||||
return self.status == GitHubStatus.CONNECTED
|
||||
|
||||
@property
|
||||
def is_error(self) -> bool:
|
||||
return self.status in {GitHubStatus.OAUTH_TIMEOUT, GitHubStatus.ERROR}
|
||||
repositories: list[NuageRepository]
|
||||
|
||||
|
||||
class ChatAssistantPublicData(BaseModel):
|
||||
chat_url: str | None = None
|
||||
class NuageRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
project_name: str = Field(default="Vibe CLI", serialization_alias="project_name")
|
||||
source: str = "vibe_code_cli"
|
||||
idempotency_key: str = Field(serialization_alias="idempotencyKey")
|
||||
message: NuageMessage
|
||||
context: NuageContext
|
||||
|
||||
|
||||
class GetChatAssistantIntegrationResponse(BaseModel):
|
||||
result: ChatAssistantPublicData
|
||||
class NuageResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
|
||||
class GetGitHubIntegrationResponse(BaseModel):
|
||||
result: GitHubPublicData
|
||||
nuage_session_id: str = Field(validation_alias="sessionId")
|
||||
nuage_web_session_id: str = Field(validation_alias="webSessionId")
|
||||
nuage_project_id: str = Field(validation_alias="projectId")
|
||||
status: str
|
||||
url: str
|
||||
|
||||
|
||||
class NuageClient:
|
||||
|
|
@ -106,16 +73,12 @@ class NuageClient:
|
|||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
workflow_id: str,
|
||||
*,
|
||||
task_queue: str | None = None,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._workflow_id = workflow_id
|
||||
self._task_queue = task_queue
|
||||
self._client = client
|
||||
self._owns_client = client is None
|
||||
self._timeout = timeout
|
||||
|
|
@ -152,74 +115,16 @@ class NuageClient:
|
|||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async def start_workflow(self, params: WorkflowParams) -> str:
|
||||
async def start(self, request: NuageRequest) -> NuageResponse:
|
||||
response = await self._http_client.post(
|
||||
f"{self._base_url}/v1/workflows/{self._workflow_id}/execute",
|
||||
f"{self._base_url}/api/v1/code/sessions",
|
||||
headers=self._headers(),
|
||||
json={
|
||||
"input": params.model_dump(mode="json"),
|
||||
"task_queue": self._task_queue,
|
||||
},
|
||||
json=request.model_dump(mode="json", by_alias=True, exclude_none=True),
|
||||
)
|
||||
if not response.is_success:
|
||||
error_msg = f"Vibe Code workflow trigger failed: {response.text}"
|
||||
raise ServiceTeleportError(error_msg)
|
||||
result = WorkflowExecuteResponse.model_validate(response.json())
|
||||
return result.execution_id
|
||||
raise ServiceTeleportError(f"Vibe Code Nuage start failed: {response.text}")
|
||||
|
||||
async def get_github_integration(self, execution_id: str) -> GitHubPublicData:
|
||||
response = await self._http_client.post(
|
||||
f"{self._base_url}/v1/workflows/executions/{execution_id}/updates",
|
||||
headers=self._headers(),
|
||||
json={"name": "get_integration", "input": {"integration_id": "github"}},
|
||||
)
|
||||
if not response.is_success:
|
||||
raise ServiceTeleportError(
|
||||
f"Failed to get GitHub integration: {response.text}"
|
||||
)
|
||||
try:
|
||||
result = GetGitHubIntegrationResponse.model_validate(response.json())
|
||||
except ValidationError as e:
|
||||
data = response.json()
|
||||
error = data.get("result", {}).get("error")
|
||||
status = data.get("result", {}).get("status")
|
||||
raise ServiceTeleportError(
|
||||
f"GitHub integration error: {error or status}"
|
||||
) from e
|
||||
return result.result
|
||||
|
||||
async def wait_for_github_connection(
|
||||
self, execution_id: str, timeout: float = 600.0, interval: float = 2.0
|
||||
) -> AsyncGenerator[GitHubPublicData, None]:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
github_data = await self.get_github_integration(execution_id)
|
||||
yield github_data
|
||||
if github_data.connected:
|
||||
return
|
||||
if github_data.is_error:
|
||||
raise ServiceTeleportError(
|
||||
github_data.error
|
||||
or f"GitHub integration failed: {github_data.status.value}"
|
||||
)
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
break
|
||||
await asyncio.sleep(min(interval, remaining))
|
||||
raise ServiceTeleportError("GitHub connection timed out")
|
||||
|
||||
async def get_chat_assistant_url(self, execution_id: str) -> str | None:
|
||||
response = await self._http_client.post(
|
||||
f"{self._base_url}/v1/workflows/executions/{execution_id}/updates",
|
||||
headers=self._headers(),
|
||||
json={
|
||||
"name": "get_integration",
|
||||
"input": {"integration_id": "chat_assistant"},
|
||||
},
|
||||
)
|
||||
if not response.is_success:
|
||||
raise ServiceTeleportError(
|
||||
f"Failed to get chat assistant integration: {response.text}"
|
||||
)
|
||||
result = GetChatAssistantIntegrationResponse.model_validate(response.json())
|
||||
return result.result.chat_url
|
||||
return NuageResponse.model_validate(response.json())
|
||||
except (ValueError, ValidationError) as e:
|
||||
raise ServiceTeleportError("Vibe Code Nuage response was invalid") from e
|
||||
|
|
|
|||
|
|
@ -6,14 +6,11 @@ from vibe.core.telemetry.send import TelemetryClient
|
|||
from vibe.core.telemetry.types import TeleportFailureStage
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.teleport.types import (
|
||||
TeleportAuthRequiredEvent,
|
||||
TeleportCheckingGitEvent,
|
||||
TeleportCompleteEvent,
|
||||
TeleportFetchingUrlEvent,
|
||||
TeleportPushingEvent,
|
||||
TeleportPushRequiredEvent,
|
||||
TeleportStartingWorkflowEvent,
|
||||
TeleportWaitingForGitHubEvent,
|
||||
TeleportYieldEvent,
|
||||
)
|
||||
|
||||
|
|
@ -29,7 +26,6 @@ def send_teleport_early_failure_telemetry(
|
|||
stage=stage,
|
||||
error_class=error_class,
|
||||
push_required=False,
|
||||
github_auth_required=False,
|
||||
nb_session_messages=nb_session_messages,
|
||||
)
|
||||
|
||||
|
|
@ -40,7 +36,6 @@ class TeleportTelemetryTracker:
|
|||
nb_session_messages: int
|
||||
stage: TeleportFailureStage
|
||||
push_required: bool = False
|
||||
github_auth_required: bool = False
|
||||
success: bool = False
|
||||
error_class: str | None = None
|
||||
|
||||
|
|
@ -55,13 +50,6 @@ class TeleportTelemetryTracker:
|
|||
self.stage = "push"
|
||||
case TeleportStartingWorkflowEvent():
|
||||
self.stage = "workflow_start"
|
||||
case TeleportWaitingForGitHubEvent():
|
||||
self.stage = "github_auth"
|
||||
case TeleportAuthRequiredEvent():
|
||||
self.github_auth_required = True
|
||||
self.stage = "github_auth"
|
||||
case TeleportFetchingUrlEvent():
|
||||
self.stage = "fetch_url"
|
||||
case TeleportCompleteEvent():
|
||||
self.success = True
|
||||
|
||||
|
|
@ -78,7 +66,6 @@ class TeleportTelemetryTracker:
|
|||
def send_success(self) -> None:
|
||||
self.telemetry_client.send_teleport_completed(
|
||||
push_required=self.push_required,
|
||||
github_auth_required=self.github_auth_required,
|
||||
nb_session_messages=self.nb_session_messages,
|
||||
)
|
||||
|
||||
|
|
@ -89,6 +76,5 @@ class TeleportTelemetryTracker:
|
|||
stage=self.stage,
|
||||
error_class=self.error_class,
|
||||
push_required=self.push_required,
|
||||
github_auth_required=self.github_auth_required,
|
||||
nb_session_messages=self.nb_session_messages,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,96 +13,61 @@ import zstandard
|
|||
from vibe.core.config import VibeConfig
|
||||
from vibe.core.session.session_logger import SessionLogger
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.teleport.experimental_nuage import (
|
||||
ExperimentalNuageClient,
|
||||
ExperimentalNuageContext,
|
||||
ExperimentalNuageDiff,
|
||||
ExperimentalNuageMessage,
|
||||
ExperimentalNuageRepository,
|
||||
ExperimentalNuageRequest,
|
||||
ExperimentalNuageTextPart,
|
||||
)
|
||||
from vibe.core.teleport.git import GitRepoInfo, GitRepository
|
||||
from vibe.core.teleport.nuage import (
|
||||
ChatAssistantParams,
|
||||
GitHubParams,
|
||||
NuageClient,
|
||||
TeleportSession,
|
||||
TextChunk,
|
||||
VibeAgent,
|
||||
WorkflowConfig,
|
||||
WorkflowIntegrations,
|
||||
WorkflowParams,
|
||||
NuageContext,
|
||||
NuageDiff,
|
||||
NuageMessage,
|
||||
NuageRepository,
|
||||
NuageRequest,
|
||||
NuageTextPart,
|
||||
)
|
||||
from vibe.core.teleport.types import (
|
||||
TeleportAuthCompleteEvent,
|
||||
TeleportAuthRequiredEvent,
|
||||
TeleportCheckingGitEvent,
|
||||
TeleportCompleteEvent,
|
||||
TeleportFetchingUrlEvent,
|
||||
TeleportPushingEvent,
|
||||
TeleportPushRequiredEvent,
|
||||
TeleportPushResponseEvent,
|
||||
TeleportSendEvent,
|
||||
TeleportStartingWorkflowEvent,
|
||||
TeleportWaitingForGitHubEvent,
|
||||
TeleportYieldEvent,
|
||||
)
|
||||
from vibe.core.utils.http import build_ssl_context
|
||||
|
||||
_DEFAULT_TELEPORT_PROMPT = "Your session has been teleported on a remote workspace. Changes of workspace has been automatically teleported. External workspace changes has NOT been teleported. Environment variables has NOT been teleported. Please continue where you left off."
|
||||
|
||||
|
||||
class TeleportService:
|
||||
def __init__(
|
||||
self,
|
||||
session_logger: SessionLogger,
|
||||
vibe_code_base_url: str,
|
||||
vibe_code_workflow_id: str,
|
||||
vibe_code_sessions_base_url: str,
|
||||
vibe_code_api_key: str,
|
||||
workdir: Path | None = None,
|
||||
*,
|
||||
vibe_code_task_queue: str | None = None,
|
||||
vibe_config: VibeConfig | None = None,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
self._session_logger = session_logger
|
||||
self._vibe_code_base_url = vibe_code_base_url
|
||||
self._vibe_code_workflow_id = vibe_code_workflow_id
|
||||
self._vibe_code_sessions_base_url = vibe_code_sessions_base_url
|
||||
self._vibe_code_api_key = vibe_code_api_key
|
||||
self._vibe_code_task_queue = vibe_code_task_queue
|
||||
self._vibe_code_project_name = (
|
||||
vibe_config.vibe_code_project_name if vibe_config else None
|
||||
)
|
||||
self._experimental_nuage_enabled = (
|
||||
vibe_config.vibe_code_experimental_nuage_enabled if vibe_config else False
|
||||
)
|
||||
self._vibe_config = vibe_config
|
||||
self._git = GitRepository(workdir)
|
||||
self._client = client
|
||||
self._owns_client = client is None
|
||||
self._timeout = timeout
|
||||
self._nuage_client_instance: NuageClient | None = None
|
||||
self._experimental_nuage_client_instance: ExperimentalNuageClient | None = None
|
||||
|
||||
async def __aenter__(self) -> TeleportService:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._timeout), verify=build_ssl_context()
|
||||
)
|
||||
if self._experimental_nuage_enabled:
|
||||
self._experimental_nuage_client_instance = ExperimentalNuageClient(
|
||||
self._vibe_code_base_url, self._vibe_code_api_key, client=self._client
|
||||
)
|
||||
else:
|
||||
self._nuage_client_instance = NuageClient(
|
||||
self._vibe_code_base_url,
|
||||
self._vibe_code_api_key,
|
||||
self._vibe_code_workflow_id,
|
||||
task_queue=self._vibe_code_task_queue,
|
||||
client=self._client,
|
||||
)
|
||||
self._nuage_client_instance = NuageClient(
|
||||
self._vibe_code_sessions_base_url,
|
||||
self._vibe_code_api_key,
|
||||
client=self._client,
|
||||
)
|
||||
await self._git.__aenter__()
|
||||
return self
|
||||
|
||||
|
|
@ -130,24 +95,12 @@ class TeleportService:
|
|||
def _nuage_client(self) -> NuageClient:
|
||||
if self._nuage_client_instance is None:
|
||||
self._nuage_client_instance = NuageClient(
|
||||
self._vibe_code_base_url,
|
||||
self._vibe_code_sessions_base_url,
|
||||
self._vibe_code_api_key,
|
||||
self._vibe_code_workflow_id,
|
||||
task_queue=self._vibe_code_task_queue,
|
||||
client=self._http_client,
|
||||
)
|
||||
return self._nuage_client_instance
|
||||
|
||||
@property
|
||||
def _experimental_nuage_client(self) -> ExperimentalNuageClient:
|
||||
if self._experimental_nuage_client_instance is None:
|
||||
self._experimental_nuage_client_instance = ExperimentalNuageClient(
|
||||
self._vibe_code_base_url,
|
||||
self._vibe_code_api_key,
|
||||
client=self._http_client,
|
||||
)
|
||||
return self._experimental_nuage_client_instance
|
||||
|
||||
async def check_supported(self) -> None:
|
||||
await self._git.get_info()
|
||||
|
||||
|
|
@ -155,23 +108,15 @@ class TeleportService:
|
|||
return await self._git.is_supported()
|
||||
|
||||
async def execute(
|
||||
self, prompt: str | None, session: TeleportSession
|
||||
self, prompt: str
|
||||
) -> AsyncGenerator[TeleportYieldEvent, TeleportSendEvent]:
|
||||
if prompt:
|
||||
lechat_user_message = prompt
|
||||
else:
|
||||
last_user_message = self._get_last_user_message(session)
|
||||
if not last_user_message:
|
||||
raise ServiceTeleportError(
|
||||
"No prompt provided and no user message found in session."
|
||||
)
|
||||
lechat_user_message = f"{last_user_message} (continue)"
|
||||
prompt = _DEFAULT_TELEPORT_PROMPT
|
||||
if not prompt:
|
||||
raise ServiceTeleportError("Teleport requires a non-empty prompt.")
|
||||
self._validate_config()
|
||||
|
||||
git_info = await self._git.get_info()
|
||||
if self._experimental_nuage_enabled:
|
||||
self._validate_experimental_nuage_git_info(git_info)
|
||||
if git_info.branch is None:
|
||||
raise ServiceTeleportError("Teleport requires a checked-out branch.")
|
||||
|
||||
yield TeleportCheckingGitEvent()
|
||||
await self._git.fetch()
|
||||
|
|
@ -196,63 +141,10 @@ class TeleportService:
|
|||
|
||||
yield TeleportStartingWorkflowEvent()
|
||||
|
||||
if self._experimental_nuage_enabled:
|
||||
result = await self._experimental_nuage_client.start(
|
||||
self._build_experimental_nuage_request(
|
||||
lechat_user_message=lechat_user_message, git_info=git_info
|
||||
)
|
||||
)
|
||||
yield TeleportCompleteEvent(url=result.url)
|
||||
return
|
||||
|
||||
execution_id = await self._nuage_client.start_workflow(
|
||||
WorkflowParams(
|
||||
prompt=prompt,
|
||||
message=[TextChunk(text=lechat_user_message)],
|
||||
config=WorkflowConfig(
|
||||
agent=VibeAgent(
|
||||
vibe_config=self._vibe_config.model_dump()
|
||||
if self._vibe_config
|
||||
else None,
|
||||
session=session,
|
||||
)
|
||||
),
|
||||
integrations=WorkflowIntegrations(
|
||||
github=self._build_github_params(git_info),
|
||||
chat_assistant=ChatAssistantParams(
|
||||
create_thread=True,
|
||||
user_message=lechat_user_message,
|
||||
project_name=self._vibe_code_project_name,
|
||||
),
|
||||
),
|
||||
)
|
||||
result = await self._nuage_client.start(
|
||||
self._build_nuage_request(prompt=prompt, git_info=git_info)
|
||||
)
|
||||
|
||||
yield TeleportWaitingForGitHubEvent()
|
||||
|
||||
auth_event_sent = False
|
||||
async for github_data in self._nuage_client.wait_for_github_connection(
|
||||
execution_id
|
||||
):
|
||||
if github_data.connected:
|
||||
break
|
||||
if not auth_event_sent and github_data.oauth_url:
|
||||
yield TeleportAuthRequiredEvent(
|
||||
oauth_url=github_data.oauth_url, message=github_data.error
|
||||
)
|
||||
auth_event_sent = True
|
||||
if github_data.error:
|
||||
yield TeleportWaitingForGitHubEvent(message=github_data.error)
|
||||
|
||||
yield TeleportAuthCompleteEvent()
|
||||
|
||||
yield TeleportFetchingUrlEvent()
|
||||
chat_url = await self._nuage_client.get_chat_assistant_url(execution_id)
|
||||
|
||||
if not chat_url:
|
||||
raise ServiceTeleportError("Chat assistant URL is not available yet")
|
||||
|
||||
yield TeleportCompleteEvent(url=chat_url)
|
||||
yield TeleportCompleteEvent(url=result.url)
|
||||
|
||||
async def _push_or_fail(self) -> None:
|
||||
if not await self._git.push_current_branch():
|
||||
|
|
@ -267,37 +159,22 @@ class TeleportService:
|
|||
)
|
||||
raise ServiceTeleportError(f"{env_var} not set.")
|
||||
|
||||
def _build_github_params(self, git_info: GitRepoInfo) -> GitHubParams:
|
||||
return GitHubParams(
|
||||
repo=f"{git_info.owner}/{git_info.repo}",
|
||||
branch=git_info.branch,
|
||||
commit=git_info.commit,
|
||||
teleported_diffs=self._compress_diff(git_info.diff or ""),
|
||||
)
|
||||
|
||||
def _build_experimental_nuage_request(
|
||||
self, *, lechat_user_message: str, git_info: GitRepoInfo
|
||||
) -> ExperimentalNuageRequest:
|
||||
if git_info.branch is None:
|
||||
raise ServiceTeleportError(
|
||||
"Experimental Nuage teleport requires a checked-out branch."
|
||||
)
|
||||
|
||||
def _build_nuage_request(
|
||||
self, *, prompt: str, git_info: GitRepoInfo
|
||||
) -> NuageRequest:
|
||||
compressed = self._compress_diff(git_info.diff)
|
||||
diff = (
|
||||
ExperimentalNuageDiff(content=compressed.decode("ascii"))
|
||||
NuageDiff(content=compressed.decode("ascii"))
|
||||
if compressed is not None
|
||||
else None
|
||||
)
|
||||
|
||||
return ExperimentalNuageRequest(
|
||||
return NuageRequest(
|
||||
idempotency_key=str(uuid4()),
|
||||
message=ExperimentalNuageMessage(
|
||||
parts=[ExperimentalNuageTextPart(text=lechat_user_message)]
|
||||
),
|
||||
context=ExperimentalNuageContext(
|
||||
message=NuageMessage(parts=[NuageTextPart(text=prompt)]),
|
||||
context=NuageContext(
|
||||
repositories=[
|
||||
ExperimentalNuageRepository(
|
||||
NuageRepository(
|
||||
repo_url=git_info.remote_url,
|
||||
branch=git_info.branch,
|
||||
commit_sha=git_info.commit,
|
||||
|
|
@ -307,14 +184,6 @@ class TeleportService:
|
|||
),
|
||||
)
|
||||
|
||||
def _validate_experimental_nuage_git_info(self, git_info: GitRepoInfo) -> None:
|
||||
if git_info.branch is not None:
|
||||
return
|
||||
|
||||
raise ServiceTeleportError(
|
||||
"Experimental Nuage teleport requires a checked-out branch."
|
||||
)
|
||||
|
||||
def _compress_diff(self, diff: str, max_size: int = 1_000_000) -> bytes | None:
|
||||
if not diff:
|
||||
return None
|
||||
|
|
@ -325,11 +194,3 @@ class TeleportService:
|
|||
"Diff too large to teleport. Please commit and push your changes first."
|
||||
)
|
||||
return encoded
|
||||
|
||||
def _get_last_user_message(self, session: TeleportSession) -> str | None:
|
||||
for msg in reversed(session.messages):
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str) and content:
|
||||
return content
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -3,15 +3,6 @@ from __future__ import annotations
|
|||
from vibe.core.types import BaseEvent
|
||||
|
||||
|
||||
class TeleportAuthRequiredEvent(BaseEvent):
|
||||
oauth_url: str
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class TeleportAuthCompleteEvent(BaseEvent):
|
||||
pass
|
||||
|
||||
|
||||
class TeleportStartingWorkflowEvent(BaseEvent):
|
||||
pass
|
||||
|
||||
|
|
@ -33,27 +24,15 @@ class TeleportPushingEvent(BaseEvent):
|
|||
pass
|
||||
|
||||
|
||||
class TeleportWaitingForGitHubEvent(BaseEvent):
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class TeleportFetchingUrlEvent(BaseEvent):
|
||||
pass
|
||||
|
||||
|
||||
class TeleportCompleteEvent(BaseEvent):
|
||||
url: str
|
||||
|
||||
|
||||
type TeleportYieldEvent = (
|
||||
TeleportAuthRequiredEvent
|
||||
| TeleportAuthCompleteEvent
|
||||
| TeleportCheckingGitEvent
|
||||
TeleportCheckingGitEvent
|
||||
| TeleportPushRequiredEvent
|
||||
| TeleportPushingEvent
|
||||
| TeleportStartingWorkflowEvent
|
||||
| TeleportWaitingForGitHubEvent
|
||||
| TeleportFetchingUrlEvent
|
||||
| TeleportCompleteEvent
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue