Skip to content

Commit a058c0d

Browse files
tharropoulostharropoulos
authored
Merge pull request #123 from tharropoulos/stream-convo
feat(streaming): add streaming conversation support
2 parents 029cbdf + c96ae12 commit a058c0d

24 files changed

Lines changed: 2696 additions & 64 deletions

‎examples/async_conversation_streaming.py‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
import asyncio
2+
import os
3+
import sys
4+
import typing
5+
import uuid
6+
7+
curr_dir = os.path.dirname(os.path.realpath(__file__))
8+
repo_root = os.path.abspath(os.path.join(curr_dir, os.pardir))
9+
sys.path.insert(1, os.path.join(repo_root, "src"))
10+
11+
import typesense
12+
13+
from typesense.types.document import (
14+
MessageChunk,
15+
SearchResponse,
16+
StreamConfigBuilder,
17+
)
18+
19+
20+
def require_env(name: str) -> str:
21+
value = os.environ.get(name)
22+
if not value:
23+
raise RuntimeError(f"Missing required environment variable: {name}")
24+
return value
25+
26+
27+
async def main() -> None:
28+
typesense_api_key = require_env("TYPESENSE_API_KEY")
29+
openai_api_key = require_env("OPENAI_API_KEY")
30+
31+
run_id = uuid.uuid4().hex
32+
history_collection = f"streaming_history_{run_id}"
33+
documents_collection = f"streaming_docs_{run_id}"
34+
model_id = f"streaming_model_{run_id}"
35+
36+
client = typesense.AsyncClient(
37+
{
38+
"api_key": typesense_api_key,
39+
"nodes": [
40+
{
41+
"host": "localhost",
42+
"port": "8108",
43+
"protocol": "http",
44+
}
45+
],
46+
"connection_timeout_seconds": 10,
47+
}
48+
)
49+
50+
try:
51+
await client.collections.create(
52+
{
53+
"name": history_collection,
54+
"fields": [
55+
{"name": "conversation_id", "type": "string"},
56+
{"name": "model_id", "type": "string"},
57+
{"name": "timestamp", "type": "int32"},
58+
{"name": "role", "type": "string", "index": False},
59+
{"name": "message", "type": "string", "index": False},
60+
],
61+
}
62+
)
63+
64+
await client.collections.create(
65+
{
66+
"name": documents_collection,
67+
"fields": [
68+
{"name": "title", "type": "string"},
69+
{
70+
"name": "embedding",
71+
"type": "float[]",
72+
"embed": {
73+
"from": ["title"],
74+
"model_config": {
75+
"model_name": "openai/text-embedding-3-small",
76+
"api_key": openai_api_key,
77+
},
78+
},
79+
},
80+
],
81+
}
82+
)
83+
84+
await client.collections[documents_collection].documents.create(
85+
{"id": "stream-1", "title": "Company profile: a developer tools firm."}
86+
)
87+
await client.collections[documents_collection].documents.create(
88+
{"id": "stream-2", "title": "Internal memo about quarterly planning."}
89+
)
90+
91+
conversation_model = await client.conversations_models.create(
92+
{
93+
"id": model_id,
94+
"model_name": "openai/gpt-4o-mini",
95+
"history_collection": history_collection,
96+
"api_key": openai_api_key,
97+
"system_prompt": (
98+
"You are an assistant for question-answering. "
99+
"Only use the provided context. Add some fluff about you Being an assistant built for Typesense Conversational Search and a brief overview of how it works"
100+
),
101+
"max_bytes": 16384,
102+
}
103+
)
104+
105+
search_parameters = {
106+
"q": "What is this document about?",
107+
"query_by": "embedding",
108+
"exclude_fields": "embedding",
109+
"prefix": False,
110+
"conversation_model_id": conversation_model["id"],
111+
}
112+
documents = client.collections[documents_collection].documents
113+
114+
# Iterate over the answer as it is generated, then read the search response.
115+
async with await documents.search_stream(search_parameters) as answer_stream:
116+
async for chunk in answer_stream:
117+
print(chunk["message"], end="", flush=True)
118+
response = await answer_stream.get_final_response()
119+
print("\n---\nFound", response["found"], "documents")
120+
121+
# Or pass callbacks to search(), which returns the search response at the end.
122+
stream_config: StreamConfigBuilder[SearchResponse[typing.Any]] = (
123+
StreamConfigBuilder()
124+
)
125+
126+
@stream_config.on_chunk
127+
def on_chunk(chunk: MessageChunk) -> None:
128+
print(chunk["message"], end="", flush=True)
129+
130+
@stream_config.on_complete
131+
def on_complete(response: SearchResponse[typing.Any]) -> None:
132+
print("\n---\nComplete response keys:", response.keys())
133+
134+
await documents.search(
135+
{
136+
**search_parameters,
137+
"conversation": True,
138+
"conversation_stream": True,
139+
"stream_config": stream_config,
140+
}
141+
)
142+
finally:
143+
await client.api_call.aclose()
144+
145+
146+
if __name__ == "__main__":
147+
asyncio.run(main())

‎examples/conversation_streaming.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import os
2+
import sys
3+
import typing
4+
import uuid
5+
6+
curr_dir = os.path.dirname(os.path.realpath(__file__))
7+
repo_root = os.path.abspath(os.path.join(curr_dir, os.pardir))
8+
sys.path.insert(1, os.path.join(repo_root, "src"))
9+
10+
import typesense
11+
12+
from typesense.types.document import (
13+
MessageChunk,
14+
SearchResponse,
15+
StreamConfigBuilder,
16+
)
17+
18+
19+
def require_env(name: str) -> str:
20+
value = os.environ.get(name)
21+
if not value:
22+
raise RuntimeError(f"Missing required environment variable: {name}")
23+
return value
24+
25+
26+
typesense_api_key = require_env("TYPESENSE_API_KEY")
27+
openai_api_key = require_env("OPENAI_API_KEY")
28+
29+
run_id = uuid.uuid4().hex
30+
history_collection = f"streaming_history_{run_id}"
31+
documents_collection = f"streaming_docs_{run_id}"
32+
model_id = f"streaming_model_{run_id}"
33+
34+
client = typesense.Client(
35+
{
36+
"api_key": typesense_api_key,
37+
"nodes": [
38+
{
39+
"host": "localhost",
40+
"port": "8108",
41+
"protocol": "http",
42+
}
43+
],
44+
"connection_timeout_seconds": 10,
45+
}
46+
)
47+
48+
client.collections.create(
49+
{
50+
"name": history_collection,
51+
"fields": [
52+
{"name": "conversation_id", "type": "string"},
53+
{"name": "model_id", "type": "string"},
54+
{"name": "timestamp", "type": "int32"},
55+
{"name": "role", "type": "string", "index": False},
56+
{"name": "message", "type": "string", "index": False},
57+
],
58+
}
59+
)
60+
61+
client.collections.create(
62+
{
63+
"name": documents_collection,
64+
"fields": [
65+
{"name": "title", "type": "string"},
66+
{
67+
"name": "embedding",
68+
"type": "float[]",
69+
"embed": {
70+
"from": ["title"],
71+
"model_config": {
72+
"model_name": "openai/text-embedding-3-small",
73+
"api_key": openai_api_key,
74+
},
75+
},
76+
},
77+
],
78+
}
79+
)
80+
81+
client.collections[documents_collection].documents.create(
82+
{"id": "stream-1", "title": "Company profile: a developer tools firm."}
83+
)
84+
client.collections[documents_collection].documents.create(
85+
{"id": "stream-2", "title": "Internal memo about a quarterly planning meeting."}
86+
)
87+
88+
conversation_model = client.conversations_models.create(
89+
{
90+
"id": model_id,
91+
"model_name": "openai/gpt-4o-mini",
92+
"history_collection": history_collection,
93+
"api_key": openai_api_key,
94+
"system_prompt": (
95+
"You are an assistant for question-answering. "
96+
"Only use the provided context. Add some fluff about you Being an assistant built for Typesense Conversational Search and a brief overview of how it works"
97+
),
98+
"max_bytes": 16384,
99+
}
100+
)
101+
102+
search_parameters = {
103+
"q": "What is this document about?",
104+
"query_by": "embedding",
105+
"exclude_fields": "embedding",
106+
"prefix": False,
107+
"conversation_model_id": conversation_model["id"],
108+
}
109+
110+
# Iterate over the answer as it is generated, then read the search response.
111+
with client.collections[documents_collection].documents.search_stream(
112+
search_parameters,
113+
) as answer_stream:
114+
for chunk in answer_stream:
115+
print(chunk["message"], end="", flush=True)
116+
response = answer_stream.get_final_response()
117+
print("\n---\nFound", response["found"], "documents")
118+
119+
# Or pass callbacks to search(), which returns the search response at the end.
120+
stream_config: StreamConfigBuilder[SearchResponse[typing.Any]] = StreamConfigBuilder()
121+
122+
123+
@stream_config.on_chunk
124+
def on_chunk(chunk: MessageChunk) -> None:
125+
print(chunk["message"], end="", flush=True)
126+
127+
128+
@stream_config.on_complete
129+
def on_complete(response: SearchResponse[typing.Any]) -> None:
130+
print("\n---\nComplete response keys:", response.keys())
131+
132+
133+
client.collections[documents_collection].documents.search(
134+
{
135+
**search_parameters,
136+
"conversation": True,
137+
"conversation_stream": True,
138+
"stream_config": stream_config,
139+
}
140+
)

0 commit comments

Comments
 (0)