Skip to content

Commit

Permalink
Add data types for pipeline routing, add route for OpenAI
Browse files Browse the repository at this point in the history
We just had a simple if-else for pipeline routing. Instead, let's add a
data type that includes the paths and optionally a `target_url`. We can
use that to add a routing for OpenAI which we'll use for testing the
Copilot provider.

In the future we can make the routings pluggable.
  • Loading branch information
jhrozek committed Jan 9, 2025
1 parent 6764c67 commit 7fa31a8
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 8 deletions.
27 changes: 26 additions & 1 deletion src/codegate/providers/copilot/mapping.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
from typing import List
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional

from pydantic import BaseModel, HttpUrl
from pydantic_settings import BaseSettings
Expand Down Expand Up @@ -43,3 +45,26 @@ class CoPilotMappings(BaseSettings):
VALIDATED_ROUTES: List[CopilotProxyRoute] = [
CopilotProxyRoute(path=path, target=target) for path, target in mappings.PROXY_ROUTES
]


class PipelineType(Enum):
FIM = "fim"
CHAT = "chat"


@dataclass
class PipelineRoute:
path: str
pipeline_type: PipelineType
target_url: Optional[str] = None


PIPELINE_ROUTES = [
PipelineRoute(
path="v1/chat/completions",
# target_url="https://api.openai.com/v1/chat/completions",
pipeline_type=PipelineType.CHAT,
),
PipelineRoute(path="v1/engines/copilot-codex/completions", pipeline_type=PipelineType.FIM),
PipelineRoute(path="chat/completions", pipeline_type=PipelineType.CHAT),
]
20 changes: 13 additions & 7 deletions src/codegate/providers/copilot/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from codegate.pipeline.factory import PipelineFactory
from codegate.pipeline.output import OutputPipelineInstance
from codegate.pipeline.secrets.manager import SecretsManager
from codegate.providers.copilot.mapping import VALIDATED_ROUTES
from codegate.providers.copilot.mapping import PIPELINE_ROUTES, VALIDATED_ROUTES, PipelineType
from codegate.providers.copilot.pipeline import (
CopilotChatPipeline,
CopilotFimPipeline,
Expand Down Expand Up @@ -153,12 +153,18 @@ def __init__(self, loop: asyncio.AbstractEventLoop):
self.context_tracking: Optional[PipelineContext] = None

def _select_pipeline(self, method: str, path: str) -> Optional[CopilotPipeline]:
if method == "POST" and path == "v1/engines/copilot-codex/completions":
logger.debug("Selected CopilotFimStrategy")
return CopilotFimPipeline(self.pipeline_factory)
if method == "POST" and path == "chat/completions":
logger.debug("Selected CopilotChatStrategy")
return CopilotChatPipeline(self.pipeline_factory)
if method != "POST":
logger.debug("Not a POST request, no pipeline selected")
return None

for route in PIPELINE_ROUTES:
if path == route.path:
if route.pipeline_type == PipelineType.FIM:
logger.debug("Selected FIM pipeline")
return CopilotFimPipeline(self.pipeline_factory)
elif route.pipeline_type == PipelineType.CHAT:
logger.debug("Selected CHAT pipeline")
return CopilotChatPipeline(self.pipeline_factory)

logger.debug("No pipeline selected")
return None
Expand Down

0 comments on commit 7fa31a8

Please sign in to comment.