viktor-hirenko commited on
Commit
5ab05da
·
1 Parent(s): a601d1a

fix: handle StopIteration in HF Inference API streaming

Browse files
Files changed (1) hide show
  1. llm_handler.py +14 -3
llm_handler.py CHANGED
@@ -112,19 +112,30 @@ class LLMHandler:
112
  stream = self.client.text_generation(
113
  prompt=full_prompt,
114
  model=self.model,
115
- max_new_tokens=1024,
116
  temperature=0.7,
117
  top_p=0.95,
118
  stream=True,
119
  details=False,
120
  )
121
 
 
 
122
  for token in stream:
123
- yield token
 
 
 
 
 
 
124
 
 
 
 
125
  except Exception as e:
126
  logger.error(f"Error streaming answer: {e}")
127
- raise
128
 
129
 
130
  def format_response(
 
112
  stream = self.client.text_generation(
113
  prompt=full_prompt,
114
  model=self.model,
115
+ max_new_tokens=512,
116
  temperature=0.7,
117
  top_p=0.95,
118
  stream=True,
119
  details=False,
120
  )
121
 
122
+ # Collect tokens and yield them
123
+ has_content = False
124
  for token in stream:
125
+ if token: # Only yield non-empty tokens
126
+ has_content = True
127
+ yield token
128
+
129
+ # If no content was generated, yield a fallback message
130
+ if not has_content:
131
+ yield "I apologize, but I couldn't generate an answer based on the provided context. Please try rephrasing your question."
132
 
133
+ except StopIteration:
134
+ # Handle empty generator
135
+ yield "I apologize, but I couldn't generate an answer. The model returned an empty response."
136
  except Exception as e:
137
  logger.error(f"Error streaming answer: {e}")
138
+ yield f"\n\n❌ Error: {str(e)}"
139
 
140
 
141
  def format_response(