Merge branch 'main' into issue-936

Author: yrobla

api/openapi.json

@@ -1478,6 +1478,16 @@
             "type": "string",
             "title": "Name"
           },
+          "config": {
+            "anyOf": [
+              {
+                "$ref": "#/components/schemas/WorkspaceConfig"
+              },
+              {
+                "type": "null"
+              }
+            ]
+          },
           "rename_to": {
             "anyOf": [
               {
@@ -1590,6 +1600,17 @@
       },
       "MuxRule": {
         "properties": {
+          "provider_name": {
+            "anyOf": [
+              {
+                "type": "string"
+              },
+              {
+                "type": "null"
+              }
+            ],
+            "title": "Provider Name"
+          },
           "provider_id": {
             "type": "string",
             "title": "Provider Id"
@@ -1842,6 +1863,27 @@
           "is_active"
         ],
         "title": "Workspace"
+      },
+      "WorkspaceConfig": {
+        "properties": {
+          "system_prompt": {
+            "type": "string",
+            "title": "System Prompt"
+          },
+          "muxing_rules": {
+            "items": {
+              "$ref": "#/components/schemas/MuxRule"
+            },
+            "type": "array",
+            "title": "Muxing Rules"
+          }
+        },
+        "type": "object",
+        "required": [
+          "system_prompt",
+          "muxing_rules"
+        ],
+        "title": "WorkspaceConfig"
       }
     }
   }

src/codegate/api/v1_models.py

@@ -5,6 +5,7 @@
 
 import pydantic
 
+import codegate.muxing.models as mux_models
 from codegate.db import models as db_models
 from codegate.extract_snippets.message_extractor import CodeSnippet
 from codegate.providers.base import BaseProvider
@@ -59,9 +60,19 @@ def from_db_workspaces(
         )
 
 
-class CreateOrRenameWorkspaceRequest(pydantic.BaseModel):
+class WorkspaceConfig(pydantic.BaseModel):
+    system_prompt: str
+
+    muxing_rules: List[mux_models.MuxRule]
+
+
+class FullWorkspace(pydantic.BaseModel):
     name: str
 
+    config: Optional[WorkspaceConfig] = None
+
+
+class CreateOrRenameWorkspaceRequest(FullWorkspace):
     # If set, rename the workspace to this name. Note that
     # the 'name' field is still required and the workspace
     # workspace must exist.

src/codegate/muxing/models.py

@@ -26,6 +26,8 @@ class MuxRule(pydantic.BaseModel):
     Represents a mux rule for a provider.
     """
 
+    # Used for exportable workspaces
+    provider_name: Optional[str] = None
     provider_id: str
     model: str
     # The type of matcher to use

src/codegate/pipeline/output.py

@@ -127,7 +127,10 @@ def _record_to_db(self) -> None:
         loop.create_task(self._db_recorder.record_context(self._input_context))
 
     async def process_stream(
-        self, stream: AsyncIterator[ModelResponse], cleanup_sensitive: bool = True, finish_stream: bool = True,
+        self,
+        stream: AsyncIterator[ModelResponse],
+        cleanup_sensitive: bool = True,
+        finish_stream: bool = True,
     ) -> AsyncIterator[ModelResponse]:
         """
         Process a stream through all pipeline steps

src/codegate/providers/copilot/provider.py

@@ -912,7 +912,9 @@ async def stream_iterator():
                 for choice in record.get("content", {}).get("choices", [])
             )
             async for record in self.output_pipeline_instance.process_stream(
-                stream_iterator(), cleanup_sensitive=False, finish_stream=finish_stream_flag,
+                stream_iterator(),
+                cleanup_sensitive=False,
+                finish_stream=finish_stream_flag,
             ):
                 chunk = record.model_dump_json(exclude_none=True, exclude_unset=True)
                 sse_data = f"data: {chunk}\n\n".encode("utf-8")

src/codegate/providers/litellmshim/generators.py

@@ -17,9 +17,9 @@ async def sse_stream_generator(stream: AsyncIterator[Any]) -> AsyncIterator[str]
                 # this might even allow us to tighten the typing of the stream
                 chunk = chunk.model_dump_json(exclude_none=True, exclude_unset=True)
             try:
-                yield f"data:{chunk}\n\n"
+                yield f"data: {chunk}\n\n"
             except Exception as e:
-                yield f"data:{str(e)}\n\n"
+                yield f"data: {str(e)}\n\n"
     except Exception as e:
         yield f"data: {str(e)}\n\n"
     finally: