WIP

Author: jhrozek

src/codegate/db/connection.py

@@ -131,7 +131,9 @@ async def record_outputs(self, outputs: List[Output]) -> Optional[Output]:
         # Just store the model respnses in the list of JSON objects.
         for output in outputs:
             full_outputs.append(output.output)
+        print("-----> FULL_OUTPUTS: ", full_outputs)
         output_db.output = json.dumps(full_outputs)
+        print("-----> DB OUTPUT: ", output_db.output)
 
         sql = text(
             """

src/codegate/providers/copilot/provider.py

@@ -575,12 +575,13 @@ def _ensure_output_processor(self) -> None:
             # Already initialized, no need to reinitialize
             return
 
-        # this is a hotfix - we shortcut before selecting the output pipeline for FIM
-        # because our FIM output pipeline is actually empty as of now. We should fix this
-        # but don't have any immediate need.
-        is_fim = self.proxy.context_tracking.metadata.get("is_fim", False)
-        if is_fim:
-            return
+        # # this is a hotfix - we shortcut before selecting the output pipeline for FIM
+        # # because our FIM output pipeline is actually empty as of now. We should fix this
+        # # but don't have any immediate need.
+        # is_fim = self.proxy.context_tracking.metadata.get("is_fim", False)
+        # if is_fim:
+        #     return
+        #
 
         logger.debug("Tracking context for pipeline processing")
         self.sse_processor = SSEProcessor()
@@ -601,16 +602,25 @@ async def _process_stream(self):
             async def stream_iterator():
                 while True:
                     incoming_record = await self.stream_queue.get()
+                    if incoming_record.get("type") == "done":
+                        break
+
                     record_content = incoming_record.get("content", {})
 
                     streaming_choices = []
                     for choice in record_content.get("choices", []):
+                        is_fim = self.proxy.context_tracking.metadata.get("is_fim", False)
+                        if is_fim:
+                            content = choice.get("text", "")
+                        else:
+                            content = choice.get("delta", {}).get("content")
+
                         streaming_choices.append(
                             StreamingChoices(
                                 finish_reason=choice.get("finish_reason", None),
                                 index=0,
                                 delta=Delta(
-                                    content=choice.get("delta", {}).get("content"), role="assistant"
+                                    content=content, role="assistant"
                                 ),
                                 logprobs=None,
                             )
@@ -624,12 +634,16 @@ async def stream_iterator():
                         model=record_content.get("model", ""),
                         object="chat.completion.chunk",
                     )
+                    print("---> YIELDING", mr)
                     yield mr
 
             async for record in self.output_pipeline_instance.process_stream(stream_iterator()):
+                print("----> RECEIVED RECORD", record)
                 chunk = record.model_dump_json(exclude_none=True, exclude_unset=True)
+                # if fim, then put the content into text
                 sse_data = f"data:{chunk}\n\n".encode("utf-8")
                 chunk_size = hex(len(sse_data))[2:] + "\r\n"
+                print("WRITING CHUNK: ", chunk)
                 self._proxy_transport_write(chunk_size.encode())
                 self._proxy_transport_write(sse_data)
                 self._proxy_transport_write(b"\r\n")
@@ -648,6 +662,7 @@ async def stream_iterator():
 
     def _process_chunk(self, chunk: bytes):
         records = self.sse_processor.process_chunk(chunk)
+        print("RECEIVED RECORDS", records)
 
         for record in records:
             if self.stream_queue is None:
@@ -658,6 +673,9 @@ def _process_chunk(self, chunk: bytes):
             self.stream_queue.put_nowait(record)
 
     def _proxy_transport_write(self, data: bytes):
+        if not self.proxy.transport or self.proxy.transport.is_closing():
+            print("TRIED TO WRITE TO A CLOSED TRANSPORT")
+            return
         self.proxy.transport.write(data)
 
     def data_received(self, data: bytes) -> None:
@@ -682,6 +700,7 @@ def data_received(self, data: bytes) -> None:
 
                     data = data[header_end + 4 :]
 
+            print("PROCESSING CHUNK: ", data)
             self._process_chunk(data)
 
     def connection_lost(self, exc: Optional[Exception]) -> None:

src/codegate/providers/copilot/streaming.py

@@ -42,8 +42,9 @@ def process_chunk(self, chunk: bytes) -> list:
                         data = json.loads(data_content)
                         records.append({"type": "data", "content": data})
                     except json.JSONDecodeError:
-                        print(f"Failed to parse JSON: {data_content}")
+                        print(f"SSEProcessor failed to parse JSON: {data_content}")
 
+        print("-----> RECORDS: ", records)
         return records
 
     def get_pending(self):