Skip to content

C
coderai
Commit fee96eb2 authored by Your Name's avatar Your Name
Browse files
Options

Combine parsers module into parser.py

parent 63c4c8a4
Hide whitespace changes
Inline Side-by-side
Showing with 135 additions and 134 deletions
codai/models/__init__.py
View file @ fee96eb2
...@@ -12,6 +12,7 @@ from .parser import ( ...@@ -12,6 +12,7 @@ from .parser import (
GrokParser, GrokParser,
PhiParser, PhiParser,
ApexBig50Parser, ApexBig50Parser,
OpenAIFormatter,
) )
from .templates import AgenticTemplateManager from .templates import AgenticTemplateManager
...@@ -29,5 +30,6 @@ __all__ = [ ...@@ -29,5 +30,6 @@ __all__ = [
'GrokParser', 'GrokParser',
'PhiParser', 'PhiParser',
'ApexBig50Parser', 'ApexBig50Parser',
'OpenAIFormatter',
'AgenticTemplateManager', 'AgenticTemplateManager',
] ]
codai/models/parser.py
View file @ fee96eb2
...@@ -17,9 +17,25 @@ import functools ...@@ -17,9 +17,25 @@ import functools
import json import json
import uuid import uuid
import re import re
import time
from difflib import get_close_matches from difflib import get_close_matches
from typing import Dict, List, Any, Optional from typing import Dict, List, Any, Optional
# Try to import litellm for response formatting
# Fall back to plain dicts if litellm is not available or doesn't export these
try:
from litellm import ModelResponse, ChatCompletionChunk, Choices, StreamingChoices, Delta, Message, Usage
LITELLM_AVAILABLE = True
except ImportError:
LITELLM_AVAILABLE = False
ModelResponse = None
ChatCompletionChunk = None
Choices = None
StreamingChoices = None
Delta = None
Message = None
Usage = None
def validate_tool_output(func): def validate_tool_output(func):
"""Decorator to validate and standardize tool call output.""" """Decorator to validate and standardize tool call output."""
...@@ -457,3 +473,120 @@ class ModelParserDispatcher: ...@@ -457,3 +473,120 @@ class ModelParserDispatcher:
"""Update the tools schema.""" """Update the tools schema."""
self.tools = tools self.tools = tools
self.parser.tools = tools self.parser.tools = tools
class OpenAIFormatter:
def __init__(self, model_name):
self.model_name = model_name
self.id = f"chatcmpl-{uuid.uuid4()}"
def format_full(self, text, prompt_tokens, completion_tokens, tool_calls=None):
"""Standard Response (Non-Streaming)"""
if LITELLM_AVAILABLE and all([ModelResponse, Choices, Message, Usage]):
try:
return ModelResponse(
id=self.id,
model=self.model_name,
object="chat.completion",
created=int(time.time()),
choices=[Choices(
finish_reason="tool_calls" if tool_calls else "stop",
index=0,
message=Message(content=text if not tool_calls else None, role="assistant", tool_calls=tool_calls)
)],
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens
)
).model_dump()
except Exception:
pass
# Fallback to plain dict if litellm fails
message = {
"role": "assistant",
"content": text if not tool_calls else None,
}
if tool_calls:
message["tool_calls"] = tool_calls
choice = {
"index": 0,
"message": message,
"finish_reason": "tool_calls" if tool_calls else "stop",
}
return {
"id": self.id,
"object": "chat.completion",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
"provider": {
"provider_name": "coderai",
"provider_id": "coderai",
},
}
def format_chunk(self, delta_text, is_final=False, usage=None):
"""Streaming Chunk (Used in a Generator)"""
if LITELLM_AVAILABLE and all([ChatCompletionChunk, StreamingChoices, Delta, (Usage if usage else True)]):
try:
return ChatCompletionChunk(
id=self.id,
model=self.model_name,
object="chat.completion.chunk",
created=int(time.time()),
choices=[StreamingChoices(
finish_reason="stop" if is_final else None,
index=0,
delta=Delta(content=delta_text, role="assistant")
)],
usage=Usage(**usage) if (usage and Usage) else None
).model_dump()
except Exception:
pass
# Fallback to plain dict if litellm fails
delta = {
"content": delta_text,
"role": "assistant",
}
choice = {
"index": 0,
"delta": delta,
"finish_reason": "stop" if is_final else None,
}
chunk = {
"id": self.id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
}
if usage and is_final:
chunk["usage"] = usage
return chunk
def format_final_chunk(self, usage: dict = None) -> dict:
"""Format the final streaming chunk with usage information."""
return self.format_chunk("", is_final=True, usage=usage)
# Backward compatibility methods
def format_litellm_full(self, text: str, prompt_tokens: int, completion_tokens: int, tool_calls=None) -> dict:
"""Backward compatibility method - calls format_full."""
return self.format_full(text, prompt_tokens, completion_tokens, tool_calls)
def format_litellm_chunk(self, delta_text: str, is_final: bool = False, usage: dict = None) -> dict:
"""Backward compatibility method - calls format_chunk."""
return self.format_chunk(delta_text, is_final, usage)
codai/models/parsers.py deleted 100644 → 0
View file @ 63c4c8a4
import time
import uuid
# Try to import litellm for response formatting
# Fall back to plain dicts if litellm is not available or doesn't export these
try:
from litellm import ModelResponse, ChatCompletionChunk, Choices, StreamingChoices, Delta, Message, Usage
LITELLM_AVAILABLE = True
except ImportError:
LITELLM_AVAILABLE = False
ModelResponse = None
ChatCompletionChunk = None
Choices = None
StreamingChoices = None
Delta = None
Message = None
Usage = None
class OpenAIFormatter:
def __init__(self, model_name):
self.model_name = model_name
self.id = f"chatcmpl-{uuid.uuid4()}"
def format_full(self, text, prompt_tokens, completion_tokens, tool_calls=None):
"""Standard Response (Non-Streaming)"""
if LITELLM_AVAILABLE and all([ModelResponse, Choices, Message, Usage]):
try:
return ModelResponse(
id=self.id,
model=self.model_name,
object="chat.completion",
created=int(time.time()),
choices=[Choices(
finish_reason="tool_calls" if tool_calls else "stop",
index=0,
message=Message(content=text if not tool_calls else None, role="assistant", tool_calls=tool_calls)
)],
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens
)
).model_dump()
except Exception:
pass
# Fallback to plain dict if litellm fails
message = {
"role": "assistant",
"content": text if not tool_calls else None,
}
if tool_calls:
message["tool_calls"] = tool_calls
choice = {
"index": 0,
"message": message,
"finish_reason": "tool_calls" if tool_calls else "stop",
}
return {
"id": self.id,
"object": "chat.completion",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
"provider": {
"provider_name": "coderai",
"provider_id": "coderai",
},
}
def format_chunk(self, delta_text, is_final=False, usage=None):
"""Streaming Chunk (Used in a Generator)"""
if LITELLM_AVAILABLE and all([ChatCompletionChunk, StreamingChoices, Delta, (Usage if usage else True)]):
try:
return ChatCompletionChunk(
id=self.id,
model=self.model_name,
object="chat.completion.chunk",
created=int(time.time()),
choices=[StreamingChoices(
finish_reason="stop" if is_final else None,
index=0,
delta=Delta(content=delta_text, role="assistant")
)],
usage=Usage(**usage) if (usage and Usage) else None
).model_dump()
except Exception:
pass
# Fallback to plain dict if litellm fails
delta = {
"content": delta_text,
"role": "assistant",
}
choice = {
"index": 0,
"delta": delta,
"finish_reason": "stop" if is_final else None,
}
chunk = {
"id": self.id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
}
if usage and is_final:
chunk["usage"] = usage
return chunk
def format_final_chunk(self, usage: dict = None) -> dict:
"""Format the final streaming chunk with usage information."""
return self.format_chunk("", is_final=True, usage=usage)
# Backward compatibility methods
def format_litellm_full(self, text: str, prompt_tokens: int, completion_tokens: int, tool_calls=None) -> dict:
"""Backward compatibility method - calls format_full."""
return self.format_full(text, prompt_tokens, completion_tokens, tool_calls)
def format_litellm_chunk(self, delta_text: str, is_final: bool = False, usage: dict = None) -> dict:
"""Backward compatibility method - calls format_chunk."""
return self.format_chunk(delta_text, is_final, usage)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment