Initial commit

This commit is contained in:
2025-07-11 13:52:19 +00:00
commit e0c489f625
362 changed files with 27286 additions and 0 deletions
View File
+23
View File
@@ -0,0 +1,23 @@
"""
TextGen service module for interacting with the Text Generation WebUI API.
"""
from .client import TextGenClient
from .models import (
ChatMessage,
ChatCompletionRequest,
ChatCompletionResponse,
CompletionRequest,
CompletionResponse,
ModelInfo,
)
__all__ = [
"TextGenClient",
"ChatMessage",
"ChatCompletionRequest",
"ChatCompletionResponse",
"CompletionRequest",
"CompletionResponse",
"ModelInfo",
]
+366
View File
@@ -0,0 +1,366 @@
"""
Client for interacting with the Text Generation WebUI API.
"""
import json
import time
import uuid
import logging
import asyncio
import aiohttp
from typing import Dict, List, Optional, Union, Any, AsyncGenerator, cast
from urllib.parse import urljoin
from .models import (
ChatMessage,
ChatCompletionRequest,
ChatCompletionResponse,
CompletionRequest,
CompletionResponse,
ModelInfo,
ModelListResponse,
ModelLoadRequest,
LogitsRequest,
)
logger = logging.getLogger(__name__)
class TextGenClient:
"""Client for interacting with the Text Generation WebUI API."""
def __init__(
self,
base_url: str = "http://textgen.localhost/v1",
api_key: Optional[str] = None,
timeout: int = 120,
):
"""
Initialize the TextGen client.
Args:
base_url: Base URL for the TextGen API
api_key: API key for authentication (optional)
timeout: Request timeout in seconds
"""
self.base_url = base_url
self.api_key = api_key
self.timeout = timeout
self._session = None
async def _ensure_session(self) -> aiohttp.ClientSession:
"""Ensure that an aiohttp session exists."""
if self._session is None or self._session.closed:
self._session = aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=self.timeout)
)
return self._session
async def close(self):
"""Close the client session."""
if self._session and not self._session.closed:
await self._session.close()
self._session = None
def _get_headers(self) -> Dict[str, str]:
"""Get headers for API requests."""
headers = {
"Content-Type": "application/json",
}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
return headers
async def _make_request(
self, method: str, endpoint: str, data: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""
Make a request to the TextGen API.
Args:
method: HTTP method (GET, POST, etc.)
endpoint: API endpoint
data: Request data
Returns:
API response as a dictionary
"""
session = await self._ensure_session()
url = urljoin(self.base_url, endpoint)
try:
async with session.request(
method=method,
url=url,
headers=self._get_headers(),
json=data,
raise_for_status=True,
) as response:
return await response.json()
except aiohttp.ClientResponseError as e:
logger.error(f"API request failed: {e.status} {e.message}")
raise
except aiohttp.ClientError as e:
logger.error(f"Request error: {str(e)}")
raise
except asyncio.TimeoutError:
logger.error(f"Request timed out after {self.timeout} seconds")
raise
async def _stream_request(
self, endpoint: str, data: Dict[str, Any]
) -> AsyncGenerator[Dict[str, Any], None]:
"""
Make a streaming request to the TextGen API.
Args:
endpoint: API endpoint
data: Request data
Yields:
Chunks of the API response
"""
session = await self._ensure_session()
url = urljoin(self.base_url, endpoint)
try:
async with session.post(
url=url,
headers=self._get_headers(),
json=data,
raise_for_status=True,
) as response:
async for line in response.content:
line = line.strip()
if not line or line == b"data: [DONE]":
continue
if line.startswith(b"data: "):
line = line[6:] # Remove "data: " prefix
try:
yield json.loads(line)
except json.JSONDecodeError:
logger.error(f"Failed to parse SSE data: {line}")
except aiohttp.ClientResponseError as e:
logger.error(f"API request failed: {e.status} {e.message}")
raise
except aiohttp.ClientError as e:
logger.error(f"Request error: {str(e)}")
raise
except asyncio.TimeoutError:
logger.error(f"Request timed out after {self.timeout} seconds")
raise
async def list_models(self) -> List[ModelInfo]:
"""
List available models.
Returns:
List of available models
"""
response = await self._make_request("GET", "internal/model/list")
model_list = ModelListResponse(**response)
return model_list.data
async def load_model(self, model_name: str, **kwargs) -> Dict[str, Any]:
"""
Load a model.
Args:
model_name: Name of the model to load
**kwargs: Additional arguments for loading the model
Returns:
Response from the API
"""
request = ModelLoadRequest(model_name=model_name, args=kwargs)
return await self._make_request("POST", "internal/model/load", request.dict())
async def chat_completion(
self, request: ChatCompletionRequest
) -> Union[ChatCompletionResponse, AsyncGenerator[Dict[str, Any], None]]:
"""
Create a chat completion.
Args:
request: Chat completion request
Returns:
Chat completion response or a stream of responses
"""
request_data = request.dict(exclude_none=True)
if request.stream:
return self._stream_request("chat/completions", request_data)
response = await self._make_request("POST", "chat/completions", request_data)
return ChatCompletionResponse(**response)
async def completion(
self, request: CompletionRequest
) -> Union[CompletionResponse, AsyncGenerator[Dict[str, Any], None]]:
"""
Create a text completion.
Args:
request: Completion request
Returns:
Completion response or a stream of responses
"""
request_data = request.dict(exclude_none=True)
if request.stream:
return self._stream_request("completions", request_data)
response = await self._make_request("POST", "completions", request_data)
return CompletionResponse(**response)
async def get_logits(self, request: LogitsRequest) -> Dict[str, Any]:
"""
Get logits for a prompt.
Args:
request: Logits request
Returns:
Logits response
"""
request_data = request.dict(exclude_none=True)
return await self._make_request("POST", "internal/logits", request_data)
async def simple_chat(
self,
messages: List[Dict[str, str]],
model: Optional[str] = None,
temperature: float = 0.7,
top_p: float = 0.9,
max_tokens: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
presence_penalty: float = 0.0,
frequency_penalty: float = 0.0,
stream: bool = False,
mode: str = "instruct",
character: Optional[str] = None,
instruction_template: Optional[str] = None,
seed: Optional[int] = None,
) -> Union[str, AsyncGenerator[str, None]]:
"""
Simple interface for chat completions.
Args:
messages: List of message dictionaries with 'role' and 'content'
model: Model to use
temperature: Sampling temperature
top_p: Nucleus sampling parameter
max_tokens: Maximum tokens to generate
stop: Stop sequences
presence_penalty: Presence penalty
frequency_penalty: Frequency penalty
stream: Whether to stream the response
mode: Mode (chat or instruct)
character: Character to use (for chat mode)
instruction_template: Instruction template (for instruct mode)
seed: Random seed for reproducibility
Returns:
Generated text or a stream of text chunks
"""
chat_messages = [ChatMessage(**msg) for msg in messages]
request = ChatCompletionRequest(
messages=chat_messages,
model=model,
temperature=temperature,
top_p=top_p,
max_tokens=max_tokens,
stop=stop,
presence_penalty=presence_penalty,
frequency_penalty=frequency_penalty,
stream=stream,
mode=mode,
character=character,
instruction_template=instruction_template,
seed=seed,
)
if stream:
async def text_stream() -> AsyncGenerator[str, None]:
stream_response = await self.chat_completion(request)
if isinstance(stream_response, AsyncGenerator):
async for chunk in stream_response:
if "choices" in chunk and chunk["choices"]:
if (
"delta" in chunk["choices"][0]
and "content" in chunk["choices"][0]["delta"]
):
yield chunk["choices"][0]["delta"]["content"]
return text_stream()
else:
response = await self.chat_completion(request)
if isinstance(response, ChatCompletionResponse):
return response.choices[0].message.content
# This should never happen due to the if/else structure, but satisfies the type checker
raise TypeError("Expected ChatCompletionResponse but got stream response")
async def simple_completion(
self,
prompt: str,
model: Optional[str] = None,
temperature: float = 0.7,
top_p: float = 0.9,
max_tokens: Optional[int] = None,
stop: Optional[Union[str, List[str]]] = None,
presence_penalty: float = 0.0,
frequency_penalty: float = 0.0,
stream: bool = False,
seed: Optional[int] = None,
) -> Union[str, AsyncGenerator[str, None]]:
"""
Simple interface for text completions.
Args:
prompt: Text prompt
model: Model to use
temperature: Sampling temperature
top_p: Nucleus sampling parameter
max_tokens: Maximum tokens to generate
stop: Stop sequences
presence_penalty: Presence penalty
frequency_penalty: Frequency penalty
stream: Whether to stream the response
seed: Random seed for reproducibility
Returns:
Generated text or a stream of text chunks
"""
request = CompletionRequest(
prompt=prompt,
model=model,
temperature=temperature,
top_p=top_p,
max_tokens=max_tokens,
stop=stop,
presence_penalty=presence_penalty,
frequency_penalty=frequency_penalty,
stream=stream,
seed=seed,
)
if stream:
async def text_stream() -> AsyncGenerator[str, None]:
stream_response = await self.completion(request)
if isinstance(stream_response, AsyncGenerator):
async for chunk in stream_response:
if "choices" in chunk and chunk["choices"]:
if "text" in chunk["choices"][0]:
yield chunk["choices"][0]["text"]
return text_stream()
else:
response = await self.completion(request)
if isinstance(response, CompletionResponse):
return response.choices[0].text
# This should never happen due to the if/else structure, but satisfies the type checker
raise TypeError("Expected CompletionResponse but got stream response")
+119
View File
@@ -0,0 +1,119 @@
"""
Example usage of the TextGen client.
"""
import asyncio
import logging
from typing import AsyncGenerator
from .client import TextGenClient
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
async def chat_example():
"""Example of using the chat completion API."""
client = TextGenClient()
try:
# Simple chat example
messages = [{"role": "user", "content": "Hello! Who are you?"}]
# Non-streaming response
logger.info("Sending chat request (non-streaming)...")
response = await client.simple_chat(
messages=messages, temperature=0.7, max_tokens=500, mode="instruct"
)
logger.info(f"Response: {response}")
# Streaming response
logger.info("Sending chat request (streaming)...")
stream_response = await client.simple_chat(
messages=messages,
temperature=0.7,
max_tokens=500,
stream=True,
mode="instruct",
)
# Check if the response is a stream
if isinstance(stream_response, AsyncGenerator):
logger.info("Streaming response:")
async for chunk in stream_response:
print(chunk, end="", flush=True)
print()
else:
logger.info(f"Expected stream but got: {stream_response}")
except Exception as e:
logger.error(f"Error: {str(e)}")
finally:
await client.close()
async def completion_example():
"""Example of using the text completion API."""
client = TextGenClient()
try:
prompt = "This is a cake recipe:\n\n1."
# Non-streaming response
logger.info("Sending completion request (non-streaming)...")
response = await client.simple_completion(
prompt=prompt, temperature=0.7, max_tokens=200
)
logger.info(f"Response: {response}")
# Streaming response
logger.info("Sending completion request (streaming)...")
stream_response = await client.simple_completion(
prompt=prompt, temperature=0.7, max_tokens=200, stream=True
)
# Check if the response is a stream
if isinstance(stream_response, AsyncGenerator):
logger.info("Streaming response:")
async for chunk in stream_response:
print(chunk, end="", flush=True)
print()
else:
logger.info(f"Expected stream but got: {stream_response}")
except Exception as e:
logger.error(f"Error: {str(e)}")
finally:
await client.close()
async def list_models_example():
"""Example of listing available models."""
client = TextGenClient()
try:
logger.info("Listing available models...")
models = await client.list_models()
for model in models:
logger.info(f"Model: {model.id}")
except Exception as e:
logger.error(f"Error: {str(e)}")
finally:
await client.close()
async def main():
"""Run all examples."""
logger.info("Running TextGen client examples")
await list_models_example()
await chat_example()
await completion_example()
if __name__ == "__main__":
asyncio.run(main())
+140
View File
@@ -0,0 +1,140 @@
"""
Data models for the TextGen API.
"""
from typing import Dict, List, Optional, Union, Any
from pydantic import BaseModel, Field
class ChatMessage(BaseModel):
"""A chat message in a conversation."""
role: str = Field(
..., description="The role of the message sender (user, assistant, system)"
)
content: str = Field(..., description="The content of the message")
name: Optional[str] = Field(None, description="The name of the sender (optional)")
class ChatCompletionRequest(BaseModel):
"""Request model for chat completions."""
messages: List[ChatMessage] = Field(
..., description="The messages in the conversation"
)
model: Optional[str] = Field(None, description="The model to use for completion")
temperature: Optional[float] = Field(0.7, description="Sampling temperature")
top_p: Optional[float] = Field(0.9, description="Nucleus sampling parameter")
max_tokens: Optional[int] = Field(
None, description="Maximum number of tokens to generate"
)
stream: Optional[bool] = Field(False, description="Whether to stream the response")
stop: Optional[Union[str, List[str]]] = Field(None, description="Stop sequences")
presence_penalty: Optional[float] = Field(0.0, description="Presence penalty")
frequency_penalty: Optional[float] = Field(0.0, description="Frequency penalty")
mode: Optional[str] = Field("chat", description="Mode (chat or instruct)")
character: Optional[str] = Field(
None, description="Character to use (for chat mode)"
)
instruction_template: Optional[str] = Field(
None, description="Instruction template (for instruct mode)"
)
seed: Optional[int] = Field(None, description="Random seed for reproducibility")
class ChatCompletionResponseChoice(BaseModel):
"""A choice in a chat completion response."""
index: int = Field(..., description="Index of the choice")
message: ChatMessage = Field(..., description="The message")
finish_reason: Optional[str] = Field(None, description="Reason for finishing")
class ChatCompletionResponse(BaseModel):
"""Response model for chat completions."""
id: str = Field(..., description="Unique identifier for the completion")
object: str = Field("chat.completion", description="Object type")
created: int = Field(..., description="Unix timestamp of creation")
model: str = Field(..., description="Model used for completion")
choices: List[ChatCompletionResponseChoice] = Field(
..., description="Completion choices"
)
usage: Dict[str, int] = Field(..., description="Token usage information")
class CompletionRequest(BaseModel):
"""Request model for text completions."""
prompt: str = Field(..., description="The prompt to complete")
model: Optional[str] = Field(None, description="The model to use for completion")
temperature: Optional[float] = Field(0.7, description="Sampling temperature")
top_p: Optional[float] = Field(0.9, description="Nucleus sampling parameter")
max_tokens: Optional[int] = Field(
None, description="Maximum number of tokens to generate"
)
stream: Optional[bool] = Field(False, description="Whether to stream the response")
stop: Optional[Union[str, List[str]]] = Field(None, description="Stop sequences")
presence_penalty: Optional[float] = Field(0.0, description="Presence penalty")
frequency_penalty: Optional[float] = Field(0.0, description="Frequency penalty")
seed: Optional[int] = Field(None, description="Random seed for reproducibility")
class CompletionResponseChoice(BaseModel):
"""A choice in a completion response."""
text: str = Field(..., description="The generated text")
index: int = Field(..., description="Index of the choice")
logprobs: Optional[Any] = Field(None, description="Log probabilities")
finish_reason: Optional[str] = Field(None, description="Reason for finishing")
class CompletionResponse(BaseModel):
"""Response model for text completions."""
id: str = Field(..., description="Unique identifier for the completion")
object: str = Field("text_completion", description="Object type")
created: int = Field(..., description="Unix timestamp of creation")
model: str = Field(..., description="Model used for completion")
choices: List[CompletionResponseChoice] = Field(
..., description="Completion choices"
)
usage: Dict[str, int] = Field(..., description="Token usage information")
class ModelInfo(BaseModel):
"""Information about a model."""
id: str = Field(..., description="Model identifier")
object: str = Field("model", description="Object type")
created: int = Field(..., description="Unix timestamp of creation")
owned_by: str = Field("user", description="Owner of the model")
permission: List[Dict[str, Any]] = Field([], description="Permissions")
root: str = Field(..., description="Root model")
parent: Optional[str] = Field(None, description="Parent model")
class LogitsRequest(BaseModel):
"""Request model for logits."""
prompt: str = Field(..., description="The prompt to get logits for")
use_samplers: bool = Field(
False, description="Whether to apply sampling parameters"
)
top_k: Optional[int] = Field(None, description="Top-k sampling parameter")
top_p: Optional[float] = Field(None, description="Top-p sampling parameter")
temperature: Optional[float] = Field(None, description="Sampling temperature")
class ModelListResponse(BaseModel):
"""Response model for model list."""
object: str = Field("list", description="Object type")
data: List[ModelInfo] = Field(..., description="List of models")
class ModelLoadRequest(BaseModel):
"""Request model for loading a model."""
model_name: str = Field(..., description="Name of the model to load")
args: Dict[str, Any] = Field({}, description="Arguments for loading the model")