Skip to content

Commit 67ad047

Browse files
committed
Added context awareness in streaming and updated DOCS
1 parent 1ac5a0d commit 67ad047

File tree

6 files changed

+55
-12
lines changed

6 files changed

+55
-12
lines changed

CHANGELOG.md

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,19 @@
22

33
All notable changes to LocalLab will be documented in this file.
44

5+
## [0.4.46] - 2024-03-14
6+
7+
### Added
8+
- Added context awareness to streaming generation
9+
- Enhanced streaming response quality with context tracking
10+
- Improved streaming response coherence by maintaining conversation history
11+
- Updated documentation with streaming context examples
12+
13+
### Fixed
14+
- Fixed streaming response formatting issues
15+
- Improved error handling in streaming generation
16+
- Enhanced token cleanup for better readability
17+
518
## [0.4.45] - 2024-03-14
619

720
### Fixed

client/python_client/locallab/client.py

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ def __init__(self, config: Union[str, LocalLabConfig, Dict[str, Any]]):
103103
self.config = config
104104
self.session: Optional[aiohttp.ClientSession] = None
105105
self.ws: Optional[websockets.WebSocketClientProtocol] = None
106+
self._stream_context = [] # Add context tracking
106107

107108
async def __aenter__(self):
108109
await self.connect()
@@ -166,20 +167,21 @@ async def _request(self, method: str, path: str, **kwargs) -> Any:
166167
await asyncio.sleep(2 ** attempt)
167168

168169
async def stream_generate(self, prompt: str, options: Optional[Union[GenerateOptions, Dict]] = None) -> AsyncGenerator[str, None]:
169-
"""Stream generated text"""
170+
"""Stream generated text with context awareness"""
170171
if isinstance(options, dict):
171172
options = GenerateOptions(**options)
172173
if options is None:
173174
options = GenerateOptions()
174175

175-
# Format data consistently
176+
# Format data with context
176177
data = {
177178
"prompt": prompt,
178179
"stream": True,
179180
"max_tokens": options.max_length,
180181
"temperature": options.temperature,
181182
"top_p": options.top_p,
182-
"model": options.model_id
183+
"model": options.model_id,
184+
"context": self._stream_context[-5:] if self._stream_context else [] # Send last 5 exchanges
183185
}
184186
# Remove None values
185187
data = {k: v for k, v in data.items() if v is not None}
@@ -195,7 +197,7 @@ async def stream_generate(self, prompt: str, options: Optional[Union[GenerateOpt
195197
yield f"Error: {error_msg}"
196198
return
197199

198-
buffer = ""
200+
current_response = ""
199201
async for line in response.content:
200202
if line:
201203
try:
@@ -217,7 +219,6 @@ async def stream_generate(self, prompt: str, options: Optional[Union[GenerateOpt
217219
# Handle different response formats
218220
text = data.get("text", data.get("response", data.get("content", "")))
219221
except json.JSONDecodeError:
220-
# If not JSON, use raw line
221222
text = line
222223

223224
if text:
@@ -235,20 +236,24 @@ async def stream_generate(self, prompt: str, options: Optional[Union[GenerateOpt
235236
.replace("user:", ""))
236237

237238
# Add space between words if needed
238-
if (buffer and
239+
if (current_response and
239240
not text.startswith(" ") and
240241
not text.startswith("\n") and
241-
not buffer.endswith(" ") and
242-
not buffer.endswith("\n")):
242+
not current_response.endswith(" ") and
243+
not current_response.endswith("\n")):
243244
text = " " + text
244245

245-
buffer += text
246+
current_response += text
246247
yield text
247248

248249
except Exception as e:
249250
logger.error(f"Error processing stream chunk: {str(e)}")
250251
yield f"\nError: Failed to process response - {str(e)}"
251252
return
253+
254+
# Update context with the full exchange
255+
self._stream_context.append({"role": "user", "content": prompt})
256+
self._stream_context.append({"role": "assistant", "content": current_response})
252257

253258
except Exception as e:
254259
logger.error(f"Stream connection error: {str(e)}")

client/python_client/setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
setup(
77
name="locallab-client",
8-
version="1.0.2",
8+
version="1.0.3",
99
author="Utkarsh Tiwari",
1010
author_email="utkarshweb2023@gmail.com",
1111
description="Official Python client for LocalLab - A local LLM server",

docs/guides/examples.md

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,31 @@ async def stream_chat():
133133
print(token, end="", flush=True)
134134
```
135135

136+
### Stream Text Generation with Context
137+
138+
The streaming generation now maintains context of the conversation for more coherent responses:
139+
140+
```python
141+
async def stream_with_context():
142+
client = LocalLabClient("http://localhost:8000")
143+
try:
144+
# First response
145+
print("Q: Tell me a story about a robot")
146+
async for token in client.stream_generate("Tell me a story about a robot"):
147+
print(token, end="", flush=True)
148+
print("\n")
149+
150+
# Follow-up question (will have context from previous response)
151+
print("Q: What happens next in the story?")
152+
async for token in client.stream_generate("What happens next in the story?"):
153+
print(token, end="", flush=True)
154+
print("\n")
155+
finally:
156+
await client.close()
157+
```
158+
159+
The client maintains a context of recent exchanges, allowing for more coherent follow-up responses. The context is automatically managed and includes up to 5 previous exchanges.
160+
136161
## Batch Processing
137162

138163
### Process Multiple Prompts

locallab/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
LocalLab - A lightweight AI inference server for running LLMs locally
33
"""
44

5-
__version__ = "0.4.45"
5+
__version__ = "0.4.46"
66

77
# Only import what's necessary initially, lazy-load the rest
88
from .logger import get_logger

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
setup(
77
name="locallab",
8-
version="0.4.45",
8+
version="0.4.46",
99
packages=find_packages(include=["locallab", "locallab.*"]),
1010
install_requires=[
1111
"fastapi>=0.95.0,<1.0.0",

0 commit comments

Comments
 (0)