Skip to content

Commit e3d39bb

Browse files
qin-ptrqin-ctx
andauthored
Fix tokens (#334)
* fix * fix: context type * Update download links to version 0.1.8 * Update download links in README_zh.md Updated download links for Mac and Windows to version 0.1.8. --------- Co-authored-by: qin-ctx <qinhaojie.exe@bytedance.com>
1 parent c3a9040 commit e3d39bb

4 files changed

Lines changed: 48 additions & 34 deletions

File tree

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ An open-source, proactive context-aware AI partner, dedicated to bringing clarit
2626

2727
🌍 Join our [Discord Group](https://discord.gg/tGj7RQ3nUR)
2828

29-
<a href="https://github.com/volcengine/MineContext/releases/download/v0.1.7/MineContext-0.1.7.dmg">🖥️ Download for Mac</a> · <a href="https://github.com/volcengine/MineContext/releases/download/v0.1.7/MineContext-0.1.7-setup.exe">💻 Download for Windows</a>
29+
<a href="https://github.com/volcengine/MineContext/releases/download/v0.1.8/MineContext-0.1.8.dmg">🖥️ Download for Mac</a> · <a href="https://github.com/volcengine/MineContext/releases/download/v0.1.8/MineContext-0.1.8-setup.exe">💻 Download for Windows</a>
3030

3131
</div>
3232

README_zh.md

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

2727
🌍 加入我们的 [Discord 社区](https://discord.gg/tGj7RQ3nUR)
2828

29-
<a href="https://github.com/volcengine/MineContext/releases/download/v0.1.7/MineContext-0.1.7.dmg">🖥️ Mac 版下载</a> · <a href="https://github.com/volcengine/MineContext/releases/download/v0.1.7/MineContext-0.1.7-setup.exe">💻 Windows 版下载</a>
29+
<a href="https://github.com/volcengine/MineContext/releases/download/v0.1.8/MineContext-0.1.8.dmg">🖥️ Mac 版下载</a> · <a href="https://github.com/volcengine/MineContext/releases/download/v0.1.8/MineContext-0.1.8-setup.exe">💻 Windows 版下载</a>
3030

3131
</div>
3232

opencontext/llm/llm_client.py

Lines changed: 38 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414
from volcenginesdkarkruntime import Ark
1515

1616
from opencontext.models.context import Vectorize
17-
from opencontext.utils.logging_utils import get_logger
1817
from opencontext.monitoring import record_processing_stage
18+
from opencontext.utils.logging_utils import get_logger
1919

2020
logger = get_logger(__name__)
2121

@@ -42,7 +42,9 @@ def __init__(self, llm_type: LLMType, config: Dict[str, Any]):
4242
if not self.api_key or not self.base_url or not self.model:
4343
raise ValueError("API key, base URL, and model must be provided")
4444
self.client = OpenAI(api_key=self.api_key, base_url=self.base_url, timeout=self.timeout)
45-
self.async_client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, timeout=self.timeout)
45+
self.async_client = AsyncOpenAI(
46+
api_key=self.api_key, base_url=self.base_url, timeout=self.timeout
47+
)
4648
if self.provider == LLMProvider.DOUBAO.value and self.llm_type == LLMType.EMBEDDING:
4749
self.client = Ark(api_key=self.api_key, base_url=self.base_url, timeout=self.timeout)
4850
self.async_client = None
@@ -268,24 +270,29 @@ def _request_embedding(self, text: str, **kwargs) -> List[float]:
268270
response = self.client.embeddings.create(model=self.model, input=[text])
269271
embedding = response.data[0].embedding
270272
else:
271-
response = self.client.multimodal_embeddings.create(model=self.model, input=[
272-
{
273-
"type": "text",
274-
"text": text
275-
}
276-
])
273+
response = self.client.multimodal_embeddings.create(
274+
model=self.model, input=[{"type": "text", "text": text}]
275+
)
277276
embedding = response.data.embedding
278277

279278
# Record token usage
280279
if hasattr(response, "usage") and response.usage:
281280
try:
282281
from opencontext.monitoring import record_token_usage
283282

283+
usage = response.usage
284+
if isinstance(usage, dict):
285+
prompt_tokens = usage.get("prompt_tokens", 0)
286+
total_tokens = usage.get("total_tokens", 0)
287+
else:
288+
prompt_tokens = usage.prompt_tokens
289+
total_tokens = usage.total_tokens
290+
284291
record_token_usage(
285292
model=self.model,
286-
prompt_tokens=response.usage.prompt_tokens,
293+
prompt_tokens=prompt_tokens,
287294
completion_tokens=0, # embedding has no completion tokens
288-
total_tokens=response.usage.total_tokens,
295+
total_tokens=total_tokens,
289296
)
290297
except ImportError:
291298
pass # Monitoring module not installed or initialized
@@ -310,24 +317,29 @@ async def _request_embedding_async(self, text: str, **kwargs) -> List[float]:
310317
response = await self.async_client.embeddings.create(model=self.model, input=[text])
311318
embedding = response.data[0].embedding
312319
else:
313-
response = self.client.multimodal_embeddings.create(model=self.model, input=[
314-
{
315-
"type": "text",
316-
"text": text
317-
}
318-
])
320+
response = self.client.multimodal_embeddings.create(
321+
model=self.model, input=[{"type": "text", "text": text}]
322+
)
319323
embedding = response.data.embedding
320324

321325
# Record token usage
322326
if hasattr(response, "usage") and response.usage:
323327
try:
324328
from opencontext.monitoring import record_token_usage
325329

330+
usage = response.usage
331+
if isinstance(usage, dict):
332+
prompt_tokens = usage.get("prompt_tokens", 0)
333+
total_tokens = usage.get("total_tokens", 0)
334+
else:
335+
prompt_tokens = usage.prompt_tokens
336+
total_tokens = usage.total_tokens
337+
326338
record_token_usage(
327339
model=self.model,
328-
prompt_tokens=response.usage.prompt_tokens,
340+
prompt_tokens=prompt_tokens,
329341
completion_tokens=0, # embedding has no completion tokens
330-
total_tokens=response.usage.total_tokens,
342+
total_tokens=total_tokens,
331343
)
332344
except ImportError:
333345
pass # Monitoring module not installed or initialized
@@ -346,20 +358,19 @@ async def _request_embedding_async(self, text: str, **kwargs) -> List[float]:
346358
logger.error(f"OpenAI API error during embedding: {e}")
347359
raise
348360

349-
350-
351361
def vectorize(self, vectorize: Vectorize, **kwargs):
352362
if vectorize.vector:
353363
return
354364
vectorize.vector = self.generate_embedding(vectorize.get_vectorize_content(), **kwargs)
355365
return
356-
366+
357367
async def vectorize_async(self, vectorize: Vectorize, **kwargs):
358368
if vectorize.vector:
359369
return
360-
vectorize.vector = await self.generate_embedding_async(vectorize.get_vectorize_content(), **kwargs)
370+
vectorize.vector = await self.generate_embedding_async(
371+
vectorize.get_vectorize_content(), **kwargs
372+
)
361373
return
362-
363374

364375
def validate(self) -> tuple[bool, str]:
365376
"""
@@ -389,7 +400,7 @@ def _extract_error_summary(error: Any) -> str:
389400
"ServiceUnavailable": "Service unavailable.",
390401
"MethodNotAllowed": "Method not allowed. Check your configuration.",
391402
}
392-
403+
393404
for code, msg in volcengine_errors.items():
394405
if code in error_msg:
395406
return msg
@@ -472,12 +483,9 @@ def _extract_error_summary(error: Any) -> str:
472483
else:
473484
return False, "Embedding model returned empty response"
474485
else:
475-
response = self.client.multimodal_embeddings.create(model=self.model, input=[
476-
{
477-
"type": "text",
478-
"text": "test"
479-
}
480-
])
486+
response = self.client.multimodal_embeddings.create(
487+
model=self.model, input=[{"type": "text", "text": "test"}]
488+
)
481489
if response.data and response.data.embedding:
482490
return True, "Embedding model validation successful"
483491
else:

opencontext/server/context_operations.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,12 @@
1414
from typing import Any, Dict, List, Optional
1515

1616
from opencontext.models.context import ProcessedContext, RawContextProperties, Vectorize
17-
from opencontext.models.enums import ContentFormat, ContextSource, ContextType
17+
from opencontext.models.enums import (
18+
ContentFormat,
19+
ContextSource,
20+
ContextType,
21+
get_context_type_options,
22+
)
1823
from opencontext.storage.global_storage import get_storage
1924
from opencontext.utils.logging_utils import get_logger
2025

@@ -217,7 +222,8 @@ def get_context_types(self) -> List[str]:
217222

218223
try:
219224
collection_names = self.storage.get_vector_collection_names()
220-
return [name for name in collection_names if name in ContextType]
225+
valid_types = get_context_type_options()
226+
return [name for name in collection_names if name in valid_types]
221227
except Exception as e:
222228
logger.exception(f"Failed to get context types: {e}")
223229
raise RuntimeError(f"Failed to get context types: {str(e)}") from e

0 commit comments

Comments
 (0)