deepwrap 0.1.0__py3-none-any.whl

deepwrap/core/__init__.py

@@ -0,0 +1,4 @@
+from .session_manager import SessionManager
+from .auth import Auth
+
+__all__ = ["SessionManager", "Auth"]

deepwrap/core/auth.py

@@ -0,0 +1,123 @@
+import time
+import requests
+import websocket
+
+from typing import Optional
+from deepwrap.config import Config
+from deepwrap.utils.browser_process import BrowserProcess
+from deepwrap.utils.port_finder import PortFinder
+from deepwrap.utils.cdp_client import CDPClient
+from deepwrap.utils.dev_tools_http import DevToolsHTTP
+from deepwrap.utils.bearer_token_extractor import BearerTokenExtractor
+
+class BrowserAuthProvider:
+    """
+    Handles browser-based authentication to obtain a Bearer token.
+     - Launches a Chromium-based browser with remote debugging enabled.
+     - Waits for the user to complete the login process.
+     - Listens for network requests in the browser to extract the Bearer token from the Authorization header.
+     - Cleans up the browser process and temporary user data directory after authentication is complete or if a timeout occurs.
+    """
+    
+    def __init__(self, login_url: str, timeout: Optional[int] = None) -> None:
+        self.login_url = login_url
+        self.timeout = timeout
+
+    def get_bearer_token(self) -> Optional[str]:
+        """
+        Perform the browser-based authentication flow to obtain a Bearer token.
+        
+        Returns:
+            The obtained Bearer token if authentication is successful, otherwise None.
+        """
+        
+        port = PortFinder.find_free_port()
+
+        browser = BrowserProcess(self.login_url, port)
+        client = None
+
+        try:
+            browser.start()
+
+            DevToolsHTTP.wait_until_ready(port)
+            websocket_url = DevToolsHTTP.get_page_websocket_url(port)
+
+            client = CDPClient(websocket_url)
+            client.enable_network()
+
+            started_at = time.time()
+
+            while True:
+                if self.timeout is not None and time.time() - started_at > self.timeout:
+                    return None
+
+                try:
+                    message = client.recv()
+                except websocket.WebSocketTimeoutException:
+                    continue
+
+                token = BearerTokenExtractor.from_cdp_message(message)
+
+                if token:
+                    return token
+
+        finally:
+            if client:
+                client.close_browser()
+                client.close()
+
+            browser.terminate()
+
+
+class Auth:
+    """
+    Handles authentication for DeepWrap interactions.
+
+    - Stores a requests.Session instance.
+    - Supports browser-based authentication through an installed Chromium browser.
+    - Captures Bearer tokens from authenticated browser requests.
+    - Applies Authorization header to the session.
+    """
+
+    def __init__(self) -> None:
+        self._is_authenticated = False
+        self.bearer_token: Optional[str] = None
+
+    def is_authed(self) -> bool:
+        """
+        Check if the session is currently authenticated.
+        
+        Returns:
+            True if authenticated, False otherwise.
+        """
+        
+        return self._is_authenticated
+
+    def browser(
+        self,
+        timeout: Optional[int] = None,
+    ) -> Optional[str]:
+        """
+        Perform browser-based authentication to obtain a Bearer token.
+        
+        Args:
+            timeout: Optional timeout in seconds for the authentication process. If None, will wait indefinitely.
+        
+        Returns:
+            The obtained Bearer token if authentication is successful, otherwise None.
+        """
+        
+        provider = BrowserAuthProvider(
+            login_url = Config.login_url,
+            timeout   = timeout,
+        )
+        
+
+        token = provider.get_bearer_token()
+
+        if not token:
+            return None
+        
+        self._is_authenticated = True
+
+        return token

deepwrap/core/session_manager.py

@@ -0,0 +1,55 @@
+import requests
+
+from typing import Optional
+from deepwrap.config import Config
+
+class SessionManager:
+    """
+    Manages HTTP sessions with appropriate headers and cookies for DeepSeek interactions.
+     - Initializes a requests.Session with default headers from Config.
+     - Allows setting of authorization header using a bearer token.
+     - Supports adding cookies to the session for domain-specific interactions.
+     - Provides get and post methods that automatically include default timeouts.
+    """
+    
+    def __init__(self, bearer_token: str, cookies: Optional[dict[str, str]] = None) -> None:
+        self.session = requests.Session()
+        self.session.headers.update(Config.headers)
+        self.session.headers["authorization"] = f"Bearer {bearer_token}"
+
+        if cookies:
+            for name, value in cookies.items():
+                self.session.cookies.set(name, value, domain = Config.base_domain)
+
+    def get(self, url: str, **kwargs) -> requests.Response:
+        """
+        Wrapper around requests.Session.get that applies default timeout and any additional kwargs.
+        
+        Args:
+            url (str): The URL to send the GET request to.
+            **kwargs: Additional keyword arguments to pass to requests.Session.get.
+            
+        Returns:
+            requests.Response: The response object resulting from the GET request.
+        """
+        
+        kwargs.setdefault("timeout", Config.default_timeout)
+        
+        return self.session.get(url, **kwargs)
+
+    def post(self, url: str, **kwargs) -> requests.Response:
+        """
+        Wrapper around requests.Session.post that applies default timeout and any additional kwargs.
+
+        Args:
+            url (str): The URL to send the POST request to.
+            **kwargs: Additional keyword arguments to pass to requests.Session.post.
+
+        Returns:
+            requests.Response: The response object resulting from the POST request.
+        """
+
+        if not kwargs.get("stream"):
+            kwargs.setdefault("timeout", Config.default_timeout)
+        
+        return self.session.post(url, **kwargs)

deepwrap/interfaces/__init__.py

@@ -0,0 +1 @@
+from .cli import DeepWrapCLI

deepwrap/interfaces/api.py

@@ -0,0 +1,373 @@
+from __future__ import annotations
+
+import hashlib
+import threading
+import uuid
+
+from dataclasses import dataclass, field
+from typing import Dict, Generator, Literal, Optional
+
+from fastapi import FastAPI, HTTPException, status
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field, field_validator
+
+from deepwrap import Client
+from deepwrap.utils.config_store import ConfigStore
+
+
+ModelName = Literal["expert", "default", "vision"]
+StreamFormat = Literal["text", "sse"]
+
+
+class ChatRequest(BaseModel):
+    message: str = Field(..., min_length=1)
+    model: ModelName = "expert"
+    token: Optional[str] = None
+    thinking: bool = True
+    search: bool = True
+    god_mode: bool = False
+    stream: bool = False
+    stream_format: StreamFormat = "text"
+    session_id: Optional[str] = None
+
+    @field_validator("message")
+    @classmethod
+    def normalize_message(cls, value: str) -> str:
+        value = value.strip()
+
+        if not value:
+            raise ValueError("Message cannot be empty.")
+
+        return value
+
+
+class ChatResponse(BaseModel):
+    model: str
+    response: str
+    session_id: Optional[str] = None
+
+
+class CreateSessionRequest(BaseModel):
+    model: ModelName = "expert"
+    token: Optional[str] = None
+    god_mode: bool = False
+
+
+class CreateSessionResponse(BaseModel):
+    session_id: str
+    model: str
+    god_mode: bool
+
+
+class DeleteSessionResponse(BaseModel):
+    ok: bool
+    session_id: str
+
+
+class HealthResponse(BaseModel):
+    ok: bool
+    app: str
+    version: str
+    token_configured: bool
+    cached_clients: int
+    active_sessions: int
+
+
+@dataclass
+class SessionRecord:
+    chat: object
+    model: str
+    token_fingerprint: str
+    god_mode: bool
+    lock: threading.Lock = field(default_factory=threading.Lock)
+
+
+class APIState:
+    def __init__(self) -> None:
+        self.store = ConfigStore()
+
+        self._clients: Dict[str, Client] = {}
+        self._sessions: Dict[str, SessionRecord] = {}
+
+        self._clients_lock = threading.Lock()
+        self._sessions_lock = threading.Lock()
+
+    @staticmethod
+    def fingerprint_token(token: str) -> str:
+        return hashlib.sha256(token.encode("utf-8")).hexdigest()
+
+    def resolve_token(self, explicit_token: Optional[str] = None) -> str:
+        if explicit_token:
+            return explicit_token.strip()
+
+        config = self.store.load()
+
+        if config.token:
+            return config.token.strip()
+
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=(
+                "No token configured. Run `deepwrap auth --browser` first "
+                "or pass `token` in the request body."
+            ),
+        )
+
+    def get_client(self, token: str) -> Client:
+        fingerprint = self.fingerprint_token(token)
+
+        with self._clients_lock:
+            client = self._clients.get(fingerprint)
+
+            if client is not None:
+                return client
+
+            try:
+                client = Client(api_key=token)
+            except Exception as exc:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail=f"Failed to initialize DeepWrap client: {exc}",
+                ) from exc
+
+            self._clients[fingerprint] = client
+
+            return client
+
+    def create_session(
+        self,
+        token: Optional[str],
+        model: str,
+        god_mode: bool,
+    ) -> CreateSessionResponse:
+        resolved_token = self.resolve_token(token)
+        fingerprint = self.fingerprint_token(resolved_token)
+        client = self.get_client(resolved_token)
+
+        try:
+            chat = client.chats.create_session(
+                model=model,
+                god_mode=god_mode,
+            )
+        except Exception as exc:
+            raise HTTPException(
+                status_code=status.HTTP_502_BAD_GATEWAY,
+                detail=f"Failed to create chat session: {exc}",
+            ) from exc
+
+        session_id = f"chat_{uuid.uuid4().hex}"
+
+        record = SessionRecord(
+            chat=chat,
+            model=model,
+            token_fingerprint=fingerprint,
+            god_mode=god_mode,
+        )
+
+        with self._sessions_lock:
+            self._sessions[session_id] = record
+
+        return CreateSessionResponse(
+            session_id=session_id,
+            model=model,
+            god_mode=god_mode,
+        )
+
+    def delete_session(self, session_id: str) -> DeleteSessionResponse:
+        with self._sessions_lock:
+            if session_id not in self._sessions:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail=f"Unknown session_id: {session_id}",
+                )
+
+            self._sessions.pop(session_id)
+
+        return DeleteSessionResponse(
+            ok=True,
+            session_id=session_id,
+        )
+
+    def get_session(self, session_id: str) -> SessionRecord:
+        with self._sessions_lock:
+            record = self._sessions.get(session_id)
+
+        if record is None:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail=f"Unknown session_id: {session_id}",
+            )
+
+        return record
+
+    def create_ephemeral_chat(
+        self,
+        token: Optional[str],
+        model: str,
+        god_mode: bool,
+    ):
+        resolved_token = self.resolve_token(token)
+        client = self.get_client(resolved_token)
+
+        try:
+            return client.chats.create_session(
+                model=model,
+                god_mode=god_mode,
+            )
+        except Exception as exc:
+            raise HTTPException(
+                status_code=status.HTTP_502_BAD_GATEWAY,
+                detail=f"Failed to create chat session: {exc}",
+            ) from exc
+
+    @property
+    def cached_clients_count(self) -> int:
+        with self._clients_lock:
+            return len(self._clients)
+
+    @property
+    def active_sessions_count(self) -> int:
+        with self._sessions_lock:
+            return len(self._sessions)
+
+
+def format_sse_chunk(chunk: str) -> str:
+    chunk = chunk.replace("\r", "")
+
+    lines = chunk.split("\n")
+
+    return "".join(f"data: {line}\n" for line in lines) + "\n"
+
+
+def create_app() -> FastAPI:
+    app = FastAPI(
+        title="DeepWrap API",
+        version="0.1.0",
+        description="Local HTTP API for DeepWrap.",
+    )
+
+    state = APIState()
+
+    @app.get("/health", response_model=HealthResponse)
+    def health() -> HealthResponse:
+        config = state.store.load()
+
+        return HealthResponse(
+            ok=True,
+            app="deepwrap",
+            version="0.1.0",
+            token_configured=bool(config.token),
+            cached_clients=state.cached_clients_count,
+            active_sessions=state.active_sessions_count,
+        )
+
+    @app.post("/sessions", response_model=CreateSessionResponse)
+    def create_session(request: CreateSessionRequest) -> CreateSessionResponse:
+        return state.create_session(
+            token=request.token,
+            model=request.model,
+            god_mode=request.god_mode,
+        )
+
+    @app.delete("/sessions/{session_id}", response_model=DeleteSessionResponse)
+    def delete_session(session_id: str) -> DeleteSessionResponse:
+        return state.delete_session(session_id)
+
+    @app.post("/chat")
+    def chat(request: ChatRequest):
+        if request.session_id:
+            record = state.get_session(request.session_id)
+
+            if record.model != request.model:
+                raise HTTPException(
+                    status_code=status.HTTP_409_CONFLICT,
+                    detail=(
+                        f"Session model mismatch. Session uses '{record.model}', "
+                        f"but request asked for '{request.model}'."
+                    ),
+                )
+
+            chat_session = record.chat
+            session_lock = record.lock
+            response_session_id = request.session_id
+
+        else:
+            chat_session = state.create_ephemeral_chat(
+                token=request.token,
+                model=request.model,
+                god_mode=request.god_mode,
+            )
+
+            session_lock = threading.Lock()
+            response_session_id = None
+
+        if request.stream:
+            def stream_text() -> Generator[str, None, None]:
+                with session_lock:
+                    try:
+                        for chunk in chat_session.respond(
+                            request.message,
+                            thinking=request.thinking,
+                            search=request.search,
+                            stream=True,
+                        ):
+                            yield chunk
+
+                    except Exception as exc:
+                        yield f"\n[DeepWrap API error] {exc}\n"
+
+            def stream_sse() -> Generator[str, None, None]:
+                with session_lock:
+                    try:
+                        for chunk in chat_session.respond(
+                            request.message,
+                            thinking=request.thinking,
+                            search=request.search,
+                            stream=True,
+                        ):
+                            yield format_sse_chunk(chunk)
+
+                        yield "event: done\ndata: [DONE]\n\n"
+
+                    except Exception as exc:
+                        yield f"event: error\ndata: {str(exc)}\n\n"
+
+            if request.stream_format == "sse":
+                return StreamingResponse(
+                    stream_sse(),
+                    media_type="text/event-stream",
+                    headers={
+                        "Cache-Control": "no-cache",
+                        "X-Accel-Buffering": "no",
+                    },
+                )
+
+            return StreamingResponse(
+                stream_text(),
+                media_type="text/plain; charset=utf-8",
+            )
+
+        with session_lock:
+            try:
+                response = chat_session.respond(
+                    request.message,
+                    thinking=request.thinking,
+                    search=request.search,
+                    stream=False,
+                )
+            except Exception as exc:
+                raise HTTPException(
+                    status_code=status.HTTP_502_BAD_GATEWAY,
+                    detail=f"Chat request failed: {exc}",
+                ) from exc
+
+        return ChatResponse(
+            model=request.model,
+            response=response,
+            session_id=response_session_id,
+        )
+
+    return app
+
+
+app = create_app()