Allows application/ndjson header & Add debug print statements for tracing execution flow in OllamaApi and EventStream

Author: Rutik Thakre

llm-ollama/src/client.rs

@@ -23,6 +23,7 @@ pub struct OllamaApi {
 
 impl OllamaApi {
     pub fn new(default_model: String) -> Self {
+        println!("OllamaApi::new");
         let base_url =
             std::env::var("GOLEM_OLLAMA_BASE_URL").unwrap_or("http://localhost:11434".to_string());
         let client = Client::builder()
@@ -36,6 +37,7 @@ impl OllamaApi {
     }
 
     pub fn send_chat(&self, params: CompletionsRequest) -> Result<CompletionsResponse, Error> {
+        println!("OllamaApi:: send_chat");
         trace!("Sending request to Ollama API: {params:?}");
 
         let mut modified_params = params;
@@ -60,6 +62,7 @@ impl OllamaApi {
     }
 
     pub fn send_chat_stream(&self, params: CompletionsRequest) -> Result<EventSource, Error> {
+        println!("OllamaApi::send_chat_stream");
         trace!("Sending request to Ollama API: {params:?}");
 
         let mut modified_params = params;

llm-ollama/src/lib.rs

@@ -20,12 +20,12 @@ mod conversions;
 struct OllamaChatStream {
     stream: RefCell<Option<EventSource>>,
     failure: Option<Error>,
-    finished: RefCell<bool>,
+    finished: RefCell<bool>, 
 }
 
 impl OllamaChatStream {
     pub fn new(stream: EventSource) -> LlmChatStream<Self> {
-        // Remove EventSource dependency and use direct HTTP response
+        println!("OllamaChatStream::new");
         LlmChatStream::new(OllamaChatStream {
             stream: RefCell::new(Some(stream)),
             failure: None,
@@ -34,6 +34,7 @@ impl OllamaChatStream {
     }
 
     pub fn failed(error: Error) -> LlmChatStream<Self> {
+        println!("OllamaChatStream::failed");
         LlmChatStream::new(OllamaChatStream {
             stream: RefCell::new(None),
             failure: Some(error),
@@ -44,25 +45,31 @@ impl OllamaChatStream {
 
 impl LlmChatStreamState for OllamaChatStream {
     fn failure(&self) -> &Option<Error> {
+        println!("OllamaChatStream::failure");
         &self.failure
     }
     fn is_finished(&self) -> bool {
+        println!("OllamaChatStream::is_finished");
         *self.finished.borrow()
     }
 
     fn set_finished(&self) {
+        println!("OllamaChatStream::set_finished");
         *self.finished.borrow_mut() = true;
     }
 
     fn stream(&self) -> Ref<Option<EventSource>> {
+        println!("OllamaChatStream::stream");
         self.stream.borrow()
     }
 
     fn stream_mut(&self) -> RefMut<Option<EventSource>> {
+        println!("OllamaChatStream::stream_mut");
         self.stream.borrow_mut()
     }
 
    fn decode_message(&self, raw: &str) -> Result<Option<StreamEvent>, String> {
+        println!("OllamaChatStream::decode_message");
         trace!("Parsing NDJSON line: {raw}");        
         let json: serde_json::Value = serde_json::from_str(raw.trim())
             .map_err(|e| format!("JSON parse error: {e}"))?;
@@ -102,6 +109,7 @@ struct OllamaComponent;
 
 impl OllamaComponent {
     fn request(client: &OllamaApi, request: CompletionsRequest) -> ChatEvent {
+        println!("OllamaComponent::request");
         match client.send_chat(request) {
             Ok(response) => process_response(response),
             Err(err) => ChatEvent::Error(err),
@@ -112,6 +120,7 @@ impl OllamaComponent {
         client: &OllamaApi,
         mut request: CompletionsRequest,
     ) -> LlmChatStream<OllamaChatStream> {
+        println!("OllamaComponent::streaming_request");
         request.stream = Some(true);
         match client.send_chat_stream(request) {
             Ok(stream) => OllamaChatStream::new(stream),
@@ -124,6 +133,7 @@ impl Guest for OllamaComponent {
     type ChatStream = LlmChatStream<OllamaChatStream>;
 
     fn send(messages: Vec<Message>, config: Config) -> ChatEvent {
+        println!("OllamaComponent::send");
         LOGGING_STATE.with_borrow_mut(|state| state.init());
 
         let client = OllamaApi::new(config.model.clone());
@@ -138,6 +148,7 @@ impl Guest for OllamaComponent {
         tool_results: Vec<(ToolCall, ToolResult)>,
         config: Config,
     ) -> ChatEvent {
+        println!("OllamaComponent::continue_");
         LOGGING_STATE.with_borrow_mut(|state| state.init());
         let client = OllamaApi::new(config.model.clone());
         match messages_to_request(messages, config.clone()) {
@@ -147,12 +158,14 @@ impl Guest for OllamaComponent {
     }
 
     fn stream(messages: Vec<Message>, config: Config) -> ChatStream {
-        ChatStream::new(Self::unwrapped_stream(messages, config))
+        println!("OllamaComponent::stream");
+        ChatStream::new(Self::unwrapped_stream(messages, config.clone()))
     }
 }
 
 impl ExtendedGuest for OllamaComponent {
     fn unwrapped_stream(messages: Vec<Message>, config: Config) -> LlmChatStream<OllamaChatStream> {
+        println!("OllamaComponent::unwrapped_stream");
         LOGGING_STATE.with_borrow_mut(|state| state.init());
 
         let client = OllamaApi::new(config.model.clone());
@@ -163,6 +176,7 @@ impl ExtendedGuest for OllamaComponent {
     }
 
     fn retry_prompt(original_messages: &[Message], partial_result: &[StreamDelta]) -> Vec<Message> {
+        println!("OllamaComponent::retry_prompt");
         let mut extended_messages = Vec::new();
 
         extended_messages.push(Message {
@@ -214,6 +228,11 @@ impl ExtendedGuest for OllamaComponent {
 
         extended_messages
     }
+
+    fn subscribe(stream: &Self::ChatStream) -> Pollable {
+        println!("OllamaComponent::subscribe");
+        stream.subscribe()
+    }
 }
 
 type DurableOllamaComponent = DurableLLM<OllamaComponent>;

llm/src/event_source/event_stream.rs

@@ -136,6 +136,7 @@ pub struct EventStream {
 impl EventStream {
     /// Initialize the EventStream with a Stream
     pub fn new(stream: InputStream) -> Self {
+        println!("EventStream::new");
         Self {
             stream: Utf8Stream::new(stream),
             buffer: String::new(),
@@ -148,21 +149,25 @@ impl EventStream {
     /// Set the last event ID of the stream. Useful for initializing the stream with a previous
     /// last event ID
     pub fn set_last_event_id(&mut self, id: impl Into<String>) {
+        println!("EventStream::set_last_event_id");
         self.last_event_id = id.into();
     }
 
     /// Get the last event ID of the stream
     pub fn last_event_id(&self) -> &str {
+        println!("EventStream::last_event_id");
         &self.last_event_id
     }
 
     pub fn subscribe(&self) -> Pollable {
+        println!("EventStream::subscribe");
         self.stream.subscribe()
     }
 
     pub fn poll_next(
         &mut self,
     ) -> Poll<Option<Result<MessageEvent, EventStreamError<StreamError>>>> {
+        println!("EventStream::poll_next");
         trace!("Polling for next event");
 
         match parse_event(&mut self.buffer, &mut self.builder) {
@@ -230,6 +235,7 @@ pub enum EventStreamError<E> {
 
 impl<E> From<Utf8StreamError<E>> for EventStreamError<E> {
     fn from(err: Utf8StreamError<E>) -> Self {
+        println!("EventStreamError::from");
         match err {
             Utf8StreamError::Utf8(err) => Self::Utf8(err),
             Utf8StreamError::Transport(err) => Self::Transport(err),
@@ -239,6 +245,7 @@ impl<E> From<Utf8StreamError<E>> for EventStreamError<E> {
 
 impl<E> From<NomError<&str>> for EventStreamError<E> {
     fn from(err: NomError<&str>) -> Self {
+        println!("EventStreamError::from");
         EventStreamError::Parser(NomError::new(err.input.to_string(), err.code))
     }
 }
@@ -248,6 +255,7 @@ where
     E: fmt::Display,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        println!("EventStreamError::fmt");
         match self {
             Self::Utf8(err) => f.write_fmt(format_args!("UTF8 error: {}", err)),
             Self::Parser(err) => f.write_fmt(format_args!("Parse error: {}", err)),
@@ -262,6 +270,7 @@ fn parse_event<E>(
     buffer: &mut String,
     builder: &mut EventBuilder,
 ) -> Result<Option<MessageEvent>, EventStreamError<E>> {
+    println!("poll_next->parse_event");
     if buffer.is_empty() {
         return Ok(None);
     }

llm/src/event_source/mod.rs

@@ -36,6 +36,7 @@ pub struct EventSource {
 impl EventSource {
     #[allow(clippy::result_large_err)]
     pub fn new(response: Response) -> Result<Self, Error> {
+        println!("EventSource::new");
         match check_response(response) {
             Ok(mut response) => {
                 let handle = unsafe {
@@ -57,11 +58,13 @@ impl EventSource {
 
     /// Close the EventSource stream and stop trying to reconnect
     pub fn close(&mut self) {
+        println!("EventSource::close");
         self.is_closed = true;
     }
 
     /// Get the current ready state
     pub fn ready_state(&self) -> ReadyState {
+        println!("EventSource::ready_state");
         if self.is_closed {
             ReadyState::Closed
         } else {
@@ -70,10 +73,12 @@ impl EventSource {
     }
 
     pub fn subscribe(&self) -> Pollable {
+        println!("EventSource::subscribe");
         self.stream.subscribe()
     }
 
     pub fn poll_next(&mut self) -> Poll<Option<Result<Event, Error>>> {
+        println!("EventSource::poll_next");
         if self.is_closed {
             return Poll::Ready(None);
         }
@@ -97,6 +102,7 @@ impl EventSource {
 
 #[allow(clippy::result_large_err)]
 fn check_response(response: Response) -> Result<Response, Error> {
+    println!("EventSource->check_response");
     match response.status() {
         StatusCode::OK => {}
         status => {
@@ -120,7 +126,7 @@ fn check_response(response: Response) -> Result<Response, Error> {
             matches!(
                 (mime_type.type_(), mime_type.subtype()),
                 (mime::TEXT, mime::EVENT_STREAM)
-            )
+            ) || mime_type.subtype().as_str().contains("ndjson") 
         })
         .unwrap_or(false)
     {
@@ -141,6 +147,7 @@ pub enum Event {
 
 impl From<MessageEvent> for Event {
     fn from(event: MessageEvent) -> Self {
+        println!("Event::from");
         Event::Message(event)
     }
 }