Skip to content

fix: tokens #333

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
qin-ctx wants to merge 2 commits into from
Closed

fix: tokens #333

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
43de005
fix
qin-ctx Jan 28, 2026
9c9155d
fix: context type
qin-ctx Dec 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 38 additions & 30 deletions opencontext/llm/llm_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
from volcenginesdkarkruntime import Ark

from opencontext.models.context import Vectorize
from opencontext.utils.logging_utils import get_logger
from opencontext.monitoring import record_processing_stage
from opencontext.utils.logging_utils import get_logger

logger = get_logger(__name__)

Expand All @@ -42,7 +42,9 @@ def __init__(self, llm_type: LLMType, config: Dict[str, Any]):
if not self.api_key or not self.base_url or not self.model:
raise ValueError("API key, base URL, and model must be provided")
self.client = OpenAI(api_key=self.api_key, base_url=self.base_url, timeout=self.timeout)
self.async_client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, timeout=self.timeout)
self.async_client = AsyncOpenAI(
api_key=self.api_key, base_url=self.base_url, timeout=self.timeout
)
if self.provider == LLMProvider.DOUBAO.value and self.llm_type == LLMType.EMBEDDING:
self.client = Ark(api_key=self.api_key, base_url=self.base_url, timeout=self.timeout)
self.async_client = None
Expand Down Expand Up @@ -268,24 +270,29 @@ def _request_embedding(self, text: str, **kwargs) -> List[float]:
response = self.client.embeddings.create(model=self.model, input=[text])
embedding = response.data[0].embedding
else:
response = self.client.multimodal_embeddings.create(model=self.model, input=[
{
"type": "text",
"text": text
}
])
response = self.client.multimodal_embeddings.create(
model=self.model, input=[{"type": "text", "text": text}]
)
embedding = response.data.embedding

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

usage = response.usage
if isinstance(usage, dict):
prompt_tokens = usage.get("prompt_tokens", 0)
total_tokens = usage.get("total_tokens", 0)
else:
prompt_tokens = usage.prompt_tokens
total_tokens = usage.total_tokens

record_token_usage(
model=self.model,
prompt_tokens=response.usage.prompt_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=0, # embedding has no completion tokens
total_tokens=response.usage.total_tokens,
total_tokens=total_tokens,
)
except ImportError:
pass # Monitoring module not installed or initialized
Expand All @@ -310,24 +317,29 @@ async def _request_embedding_async(self, text: str, **kwargs) -> List[float]:
response = await self.async_client.embeddings.create(model=self.model, input=[text])
embedding = response.data[0].embedding
else:
response = self.client.multimodal_embeddings.create(model=self.model, input=[
{
"type": "text",
"text": text
}
])
response = self.client.multimodal_embeddings.create(
model=self.model, input=[{"type": "text", "text": text}]
)
embedding = response.data.embedding

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

usage = response.usage
if isinstance(usage, dict):
prompt_tokens = usage.get("prompt_tokens", 0)
total_tokens = usage.get("total_tokens", 0)
else:
prompt_tokens = usage.prompt_tokens
total_tokens = usage.total_tokens

record_token_usage(
model=self.model,
prompt_tokens=response.usage.prompt_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=0, # embedding has no completion tokens
total_tokens=response.usage.total_tokens,
total_tokens=total_tokens,
)
except ImportError:
pass # Monitoring module not installed or initialized
Expand All @@ -346,20 +358,19 @@ async def _request_embedding_async(self, text: str, **kwargs) -> List[float]:
logger.error(f"OpenAI API error during embedding: {e}")
raise



def vectorize(self, vectorize: Vectorize, **kwargs):
if vectorize.vector:
return
vectorize.vector = self.generate_embedding(vectorize.get_vectorize_content(), **kwargs)
return

async def vectorize_async(self, vectorize: Vectorize, **kwargs):
if vectorize.vector:
return
vectorize.vector = await self.generate_embedding_async(vectorize.get_vectorize_content(), **kwargs)
vectorize.vector = await self.generate_embedding_async(
vectorize.get_vectorize_content(), **kwargs
)
return


def validate(self) -> tuple[bool, str]:
"""
Expand Down Expand Up @@ -389,7 +400,7 @@ def _extract_error_summary(error: Any) -> str:
"ServiceUnavailable": "Service unavailable.",
"MethodNotAllowed": "Method not allowed. Check your configuration.",
}

for code, msg in volcengine_errors.items():
if code in error_msg:
return msg
Expand Down Expand Up @@ -472,12 +483,9 @@ def _extract_error_summary(error: Any) -> str:
else:
return False, "Embedding model returned empty response"
else:
response = self.client.multimodal_embeddings.create(model=self.model, input=[
{
"type": "text",
"text": "test"
}
])
response = self.client.multimodal_embeddings.create(
model=self.model, input=[{"type": "text", "text": "test"}]
)
if response.data and response.data.embedding:
return True, "Embedding model validation successful"
else:
Expand Down
10 changes: 8 additions & 2 deletions opencontext/server/context_operations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,12 @@
from typing import Any, Dict, List, Optional

from opencontext.models.context import ProcessedContext, RawContextProperties, Vectorize
from opencontext.models.enums import ContentFormat, ContextSource, ContextType
from opencontext.models.enums import (
ContentFormat,
ContextSource,
ContextType,
get_context_type_options,
)
from opencontext.storage.global_storage import get_storage
from opencontext.utils.logging_utils import get_logger

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

try:
collection_names = self.storage.get_vector_collection_names()
return [name for name in collection_names if name in ContextType]
valid_types = get_context_type_options()
return [name for name in collection_names if name in valid_types]
except Exception as e:
logger.exception(f"Failed to get context types: {e}")
raise RuntimeError(f"Failed to get context types: {str(e)}") from e
Loading