@@ -47,13 +47,18 @@ def __init__(self, ollama_response: AsyncIterator[ChatResponse]):
47
47
self ._aiter = ollama_response .__aiter__ ()
48
48
49
49
@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
+ """
52
54
datetime_obj = datetime .fromisoformat (chunk_created_at )
53
55
return int (datetime_obj .timestamp ())
54
56
55
57
@classmethod
56
58
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
+ """
57
62
finish_reason = None
58
63
role = "assistant"
59
64
if is_chunk_done :
@@ -63,14 +68,20 @@ def _get_finish_reason_assistant(cls, is_chunk_done: bool) -> Tuple[str, Optiona
63
68
64
69
@classmethod
65
70
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
+ """
66
76
timestamp_str = str (timestamp_seconds )
67
77
return timestamp_str [:9 ]
68
78
69
79
@classmethod
70
80
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
+ """
72
84
timestamp_seconds = cls ._transform_to_int_secs (chunk .created_at )
73
- # Get role and finish reason
74
85
role , finish_reason = cls ._get_finish_reason_assistant (chunk .done )
75
86
chat_id = cls ._get_chat_id_from_timestamp (timestamp_seconds )
76
87
@@ -95,9 +106,7 @@ def normalize_fim_chunk(cls, chunk: GenerateResponse) -> Dict:
95
106
"""
96
107
Transform an ollama generation chunk to an OpenAI one
97
108
"""
98
- # Convert the datetime object to a timestamp in seconds
99
109
timestamp_seconds = cls ._transform_to_int_secs (chunk .created_at )
100
- # Get role and finish reason
101
110
_ , finish_reason = cls ._get_finish_reason_assistant (chunk .done )
102
111
chat_id = cls ._get_chat_id_from_timestamp (timestamp_seconds )
103
112
0 commit comments