Skip to content

Commit 60acedc

Browse files
Adressed comments from review
1 parent 095bcbb commit 60acedc

File tree

2 files changed

+18
-9
lines changed

2 files changed

+18
-9
lines changed

src/codegate/muxing/adapter.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import uuid
44
from abc import ABC, abstractmethod
55
from typing import Callable, Dict, Union
6+
from urllib.parse import urljoin
67

78
import structlog
89
from fastapi.responses import JSONResponse, StreamingResponse
@@ -31,13 +32,12 @@ class BodyAdapter:
3132

3233
def _get_provider_formatted_url(self, model_route: rulematcher.ModelRoute) -> str:
3334
"""Get the provider formatted URL to use in base_url. Note this value comes from DB"""
34-
base_endpoint = model_route.endpoint.endpoint.rstrip("/")
3535
if model_route.endpoint.provider_type in [
3636
db_models.ProviderType.openai,
3737
db_models.ProviderType.openrouter,
3838
]:
39-
return f"{base_endpoint}/v1"
40-
return base_endpoint
39+
return urljoin(model_route.endpoint.endpoint, "/v1")
40+
return model_route.endpoint.endpoint
4141

4242
def set_destination_info(self, model_route: rulematcher.ModelRoute, data: dict) -> dict:
4343
"""Set the destination provider info."""

src/codegate/providers/ollama/adapter.py

+15-6
Original file line numberDiff line numberDiff line change
@@ -47,13 +47,18 @@ def __init__(self, ollama_response: AsyncIterator[ChatResponse]):
4747
self._aiter = ollama_response.__aiter__()
4848

4949
@classmethod
50-
def _transform_to_int_secs(cls, chunk_created_at) -> int:
51-
# Convert the datetime object to a timestamp in seconds
50+
def _transform_to_int_secs(cls, chunk_created_at: str) -> int:
51+
"""
52+
Convert the datetime to a timestamp in seconds.
53+
"""
5254
datetime_obj = datetime.fromisoformat(chunk_created_at)
5355
return int(datetime_obj.timestamp())
5456

5557
@classmethod
5658
def _get_finish_reason_assistant(cls, is_chunk_done: bool) -> Tuple[str, Optional[str]]:
59+
"""
60+
Get the role and finish reason for the assistant based on the chunk done status.
61+
"""
5762
finish_reason = None
5863
role = "assistant"
5964
if is_chunk_done:
@@ -63,14 +68,20 @@ def _get_finish_reason_assistant(cls, is_chunk_done: bool) -> Tuple[str, Optiona
6368

6469
@classmethod
6570
def _get_chat_id_from_timestamp(cls, timestamp_seconds: int) -> str:
71+
"""
72+
Getting a string representation of the timestamp in seconds used as the chat id.
73+
74+
This needs to be done so that all chunks of a chat have the same id.
75+
"""
6676
timestamp_str = str(timestamp_seconds)
6777
return timestamp_str[:9]
6878

6979
@classmethod
7080
def normalize_chat_chunk(cls, chunk: ChatResponse) -> ModelResponse:
71-
# Convert the datetime object to a timestamp in seconds
81+
"""
82+
Transform an ollama chat chunk to an OpenAI one
83+
"""
7284
timestamp_seconds = cls._transform_to_int_secs(chunk.created_at)
73-
# Get role and finish reason
7485
role, finish_reason = cls._get_finish_reason_assistant(chunk.done)
7586
chat_id = cls._get_chat_id_from_timestamp(timestamp_seconds)
7687

@@ -95,9 +106,7 @@ def normalize_fim_chunk(cls, chunk: GenerateResponse) -> Dict:
95106
"""
96107
Transform an ollama generation chunk to an OpenAI one
97108
"""
98-
# Convert the datetime object to a timestamp in seconds
99109
timestamp_seconds = cls._transform_to_int_secs(chunk.created_at)
100-
# Get role and finish reason
101110
_, finish_reason = cls._get_finish_reason_assistant(chunk.done)
102111
chat_id = cls._get_chat_id_from_timestamp(timestamp_seconds)
103112

0 commit comments

Comments
 (0)