Skip to content

Commit 3b01383

Browse files
wangyi01wangyi01
committed
python sdk
1 parent 0eeae3f commit 3b01383

File tree

4 files changed

+341
-103
lines changed

4 files changed

+341
-103
lines changed

‎src/bohrium/_base_client.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,10 @@ def _request(
193193
"params": merged_params,
194194
}
195195

196+
# 处理超时参数
197+
if "timeout" in kwargs:
198+
request_kwargs["timeout"] = kwargs["timeout"]
199+
196200
if json is not None:
197201
request_kwargs["json"] = json
198202
request_kwargs["headers"] = merged_headers

‎src/bohrium/resources/sigma_search/sigma_search.py‎

Lines changed: 51 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from typing import Optional, List, Dict, Any, Union, Iterator
33
from pprint import pprint
44
import json
5+
import httpx
56

67
from ..._resource import AsyncAPIResource, SyncAPIResource
78
from ..._response import APIResponse
@@ -41,7 +42,7 @@ def create_session(
4142
if kwargs:
4243
data.update(kwargs)
4344

44-
response = self._client.post("/openapi/v1/sigma-search/api/v1/ai_search/sessions", json=data)
45+
response = self._client.post("/openapi/v1/sigma-search/api/v2/ai_search/sessions", json=data)
4546
log.info(response.json())
4647
return APIResponse(response).json.get("data")
4748

@@ -53,7 +54,7 @@ def get_session(
5354
"""获取会话详情"""
5455
log.info(f"getting sigma search session: {uuid}")
5556

56-
response = self._client.get(f"/openapi/v1/sigma-search/api/v1/ai_search/sessions/{uuid}")
57+
response = self._client.get(f"/openapi/v1/sigma-search/api/v1/ai_search/sessions_extended/{uuid}")
5758
log.info(response.json())
5859
return APIResponse(response).json.get("data")
5960

@@ -85,24 +86,54 @@ def get_summary_stream(
8586
"""获取总结流式输出"""
8687
log.info(f"getting summary stream for query: {query_id}")
8788

88-
response = self._client.get(
89-
f"/openapi/v1/sigma-search/api/v1/ai_search/questions/{query_id}/stream",
90-
stream=True
91-
)
92-
93-
# 处理流式响应
94-
for line in response.iter_lines():
95-
if line:
96-
line_str = line.decode('utf-8')
97-
if line_str.startswith('event:data'):
98-
continue
99-
elif line_str.startswith('data:'):
100-
try:
101-
data = json.loads(line_str[5:]) # 去掉 'data:' 前缀
102-
yield data
103-
except json.JSONDecodeError:
104-
log.warning(f"Failed to parse JSON: {line_str}")
105-
continue
89+
try:
90+
# 使用专门的流式HTTP客户端
91+
import httpx
92+
93+
# 创建专门的流式客户端，禁用缓冲
94+
stream_client = httpx.Client(
95+
timeout=httpx.Timeout(timeout=600.0, connect=10.0),
96+
limits=httpx.Limits(max_connections=1, max_keepalive_connections=1)
97+
)
98+
99+
# 构建完整URL
100+
url = f"{self._client._base_url}/openapi/v1/sigma-search/api/v1/ai_search/questions/{query_id}/stream"
101+
102+
# 添加access key参数
103+
params = {"accessKey": self._client.access_key}
104+
105+
# 使用专门的流式请求头
106+
headers = {
107+
"Accept": "*/*",
108+
"Connection": "keep-alive",
109+
"Cache-Control": "no-cache"
110+
}
111+
112+
log.info("开始流式请求...")
113+
114+
# 使用stream方法进行真正的流式请求
115+
with stream_client.stream(
116+
"GET",
117+
url,
118+
params=params,
119+
headers=headers
120+
) as response:
121+
log.info(f"流式响应状态: {response.status_code}")
122+
123+
if response.status_code != 200:
124+
log.error(f"流式请求失败: {response.status_code}")
125+
return
126+
127+
# 逐行读取流式数据
128+
for line in response.iter_lines():
129+
if line:
130+
log.debug(f"收到流式数据: {line[:100]}...")
131+
yield line.encode('utf-8')
132+
133+
stream_client.close()
134+
except Exception as e:
135+
log.error(f"Stream error: {e}")
136+
return
106137

107138
def get_summary_content(
108139
self,

‎tests/sigma_search.py‎

Lines changed: 83 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -79,31 +79,31 @@ def main():
7979

8080
print("\n" + "="*50 + "\n")
8181

82-
# 3. 获取相关文献
83-
print("3. 获取相关文献:")
84-
print("-" * 50)
85-
if query_id:
86-
try:
87-
papers_result = bohrium.sigma_search.get_papers(
88-
query_id=query_id,
89-
sort="RelevanceScore"
90-
)
91-
print("文献获取成功！")
92-
papers = papers_result.get('list', [])
93-
print(f"找到 {len(papers)} 篇相关文献")
82+
# # 3. 获取相关文献
83+
# print("3. 获取相关文献:")
84+
# print("-" * 50)
85+
# if query_id:
86+
# try:
87+
# papers_result = bohrium.sigma_search.get_papers(
88+
# query_id=query_id,
89+
# sort="RelevanceScore"
90+
# )
91+
# print("文献获取成功！")
92+
# papers = papers_result.get('list', [])
93+
# print(f"找到 {len(papers)} 篇相关文献")
9494

95-
if papers:
96-
paper = papers[0]
97-
print(f"第一篇文献标题: {paper.get('title')}")
98-
print(f"第一篇文献作者: {', '.join(paper.get('author', [])[:3])}")
99-
print(f"第一篇文献期刊: {paper.get('journal')}")
100-
print(f"第一篇文献摘要: {paper.get('abstract', '')[:100]}...")
101-
except Exception as e:
102-
print(f"获取文献失败: {e}")
103-
else:
104-
print("跳过文献获取测试（缺少query_id）")
105-
106-
print("\n" + "="*50 + "\n")
95+
# if papers:
96+
# paper = papers[0]
97+
# print(f"第一篇文献标题: {paper.get('title')}")
98+
# print(f"第一篇文献作者: {', '.join(paper.get('author', [])[:3])}")
99+
# print(f"第一篇文献期刊: {paper.get('journal')}")
100+
# print(f"第一篇文献摘要: {paper.get('abstract', '')[:100]}...")
101+
# except Exception as e:
102+
# print(f"获取文献失败: {e}")
103+
# else:
104+
# print("跳过文献获取测试（缺少query_id）")
105+
106+
# print("\n" + "="*50 + "\n")
107107

108108
# 4. 获取总结内容
109109
print("4. 获取总结内容:")
@@ -135,70 +135,70 @@ def main():
135135

136136
print("\n" + "="*50 + "\n")
137137

138-
# 5. 追问测试
139-
print("5. 追问测试:")
140-
print("-" * 50)
141-
try:
142-
follow_up_result = bohrium.sigma_search.follow_up_question(
143-
session_uuid=session_uuid,
144-
query="Transformer在音频处理中的应用"
145-
)
146-
print("追问成功！")
147-
print(f"新问题ID: {follow_up_result.get('id')}")
148-
print(f"新问题查询: {follow_up_result.get('query')}")
149-
print(f"最后答案ID: {follow_up_result.get('lastAnswerID')}")
150-
except Exception as e:
151-
print(f"追问失败: {e}")
152-
153-
print("\n" + "="*50 + "\n")
154-
155-
# 6. 获取搜索历史
156-
print("6. 获取搜索历史:")
157-
print("-" * 50)
158-
try:
159-
history_result = bohrium.sigma_search.get_search_history()
160-
print("搜索历史获取成功！")
161-
sessions = history_result.get('sessions', [])
162-
print(f"历史会话数量: {len(sessions)}")
138+
# # 5. 追问测试
139+
# print("5. 追问测试:")
140+
# print("-" * 50)
141+
# try:
142+
# follow_up_result = bohrium.sigma_search.follow_up_question(
143+
# session_uuid=session_uuid,
144+
# query="Transformer在音频处理中的应用"
145+
# )
146+
# print("追问成功！")
147+
# print(f"新问题ID: {follow_up_result.get('id')}")
148+
# print(f"新问题查询: {follow_up_result.get('query')}")
149+
# print(f"最后答案ID: {follow_up_result.get('lastAnswerID')}")
150+
# except Exception as e:
151+
# print(f"追问失败: {e}")
152+
153+
# print("\n" + "="*50 + "\n")
154+
155+
# # 6. 获取搜索历史
156+
# print("6. 获取搜索历史:")
157+
# print("-" * 50)
158+
# try:
159+
# history_result = bohrium.sigma_search.get_search_history()
160+
# print("搜索历史获取成功！")
161+
# sessions = history_result.get('sessions', [])
162+
# print(f"历史会话数量: {len(sessions)}")
163163

164-
for i, session in enumerate(sessions[:3]):
165-
print(f" 会话{i+1}: {session.get('title')} ({session.get('uuid')})")
166-
print(f" 创建时间: {session.get('createTime')}")
167-
print(f" 状态: {session.get('status')}")
168-
except Exception as e:
169-
print(f"获取搜索历史失败: {e}")
170-
171-
print("\n" + "="*50 + "\n")
172-
173-
# 7. 使用类型定义方式
174-
print("7. 使用类型定义方式:")
175-
print("-" * 50)
176-
try:
177-
# 创建会话请求
178-
session_request = CreateSessionRequest(
179-
query="深度学习在医疗诊断中的应用",
180-
model="qwen",
181-
discipline="LS", # 生命科学
182-
resource_id_list=[]
183-
)
164+
# for i, session in enumerate(sessions[:3]):
165+
# print(f" 会话{i+1}: {session.get('title')} ({session.get('uuid')})")
166+
# print(f" 创建时间: {session.get('createTime')}")
167+
# print(f" 状态: {session.get('status')}")
168+
# except Exception as e:
169+
# print(f"获取搜索历史失败: {e}")
170+
171+
# print("\n" + "="*50 + "\n")
172+
173+
# # 7. 使用类型定义方式
174+
# print("7. 使用类型定义方式:")
175+
# print("-" * 50)
176+
# try:
177+
# # 创建会话请求
178+
# session_request = CreateSessionRequest(
179+
# query="深度学习在医疗诊断中的应用",
180+
# model="qwen",
181+
# discipline="LS", # 生命科学
182+
# resource_id_list=[]
183+
# )
184184

185-
session_result2 = bohrium.sigma_search.search_with_request(session_request)
186-
print("使用类型定义创建会话成功！")
187-
print(f"会话UUID: {session_result2.get('uuid')}")
188-
print(f"会话标题: {session_result2.get('title')}")
185+
# session_result2 = bohrium.sigma_search.search_with_request(session_request)
186+
# print("使用类型定义创建会话成功！")
187+
# print(f"会话UUID: {session_result2.get('uuid')}")
188+
# print(f"会话标题: {session_result2.get('title')}")
189189

190-
# 追问请求
191-
follow_up_request = FollowUpRequest(
192-
session_uuid=session_result2.get('uuid'),
193-
query="深度学习在医学影像分析中的具体应用"
194-
)
190+
# # 追问请求
191+
# follow_up_request = FollowUpRequest(
192+
# session_uuid=session_result2.get('uuid'),
193+
# query="深度学习在医学影像分析中的具体应用"
194+
# )
195195

196-
follow_up_result2 = bohrium.sigma_search.search_with_request(follow_up_request)
197-
print("使用类型定义追问成功！")
198-
print(f"新问题ID: {follow_up_result2.get('id')}")
196+
# follow_up_result2 = bohrium.sigma_search.search_with_request(follow_up_request)
197+
# print("使用类型定义追问成功！")
198+
# print(f"新问题ID: {follow_up_result2.get('id')}")
199199

200-
except Exception as e:
201-
print(f"使用类型定义方式失败: {e}")
200+
# except Exception as e:
201+
# print(f"使用类型定义方式失败: {e}")
202202

203203
print("\n=== 测试完成 ===")
204204

0 commit comments

Comments
 (0)