feat: Add max_tokens support to .extract(...)
- Dominant language
- Python
- Stars
- 38.6k
- Forks
- 2.7k
- Avg merge
- 3d 15h
- Merged PRs (30d)
- 3
Description
I'm using a closed network OpenAI API implementation to access an LLM. I'm using the following bit of code to create an OpenAI client, since base_url is not available yet, as mentioned in https://github.com/google/langextract/issues/53
```python
from langextract.inference import OpenAILanguageModel
from openai import OpenAI
class LangModel(OpenAILanguageModel):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._client = OpenAI(
api_key="yo",
base_url="/api/openai/v1",
)
```
The OpenAI client has a max_tokens parameter, but I don't see a way to pass max_tokens via `.extract()`. For example,
```python
# 1. Define the prompt and extraction rules
prompt = textwrap.dedent("""\
Extract characters, emotions, and relationships in order of appearance.
Use exact text for extractions. Do not paraphrase or overlap entities.
Provide meaningful attributes for each entity to add context.""")
# 2. Provide a high-quality example to guide the model
examples = [
lx.data.ExampleData(
text="ROMEO. But soft! What light through yonder window breaks? It is the east, and Juliet is the sun.",
extractions=[
lx.data.Extraction(
extraction_class="character",
extraction_text="ROMEO",
attributes={"emotional_state": "wonder"},
),
lx.data.Extraction(
extraction_class="emotion",
extraction_text="But soft!",
attributes={"feeling": "gentle awe"},
),
lx.data.Extraction(
extraction_class="relationship",
extraction_text="Juliet is the sun",
attributes={"type": "metaphor"},
),
],
)
]
# The input text to be processed
input_text = "Lady Juliet gazed longingly at the stars, her heart aching for Romeo"
# Run the extraction
result = lx.extract(
text_or_documents=input_text,
prompt_description=prompt,
examples=examples,
language_model_type=LangModel,
model_id="general",
model_url="https://llm.t3/dev/api/openai/v1",
api_key="yo",
language_model_params={"max_output_tokens": 1024}, # hopeful that this would have worked
)
```
raises an error because max_tokens is None. Walking through the trace brought me to
```python
def _process_single_prompt(self, prompt: str, config: dict) -> ScoredOutput:
"""Process a single prompt and return a ScoredOutput."""
try:
# Prepare the system message for structured output
system_message = ''
if self.format_type == data.FormatType.JSON:
system_message = (
'You are a helpful assistant that responds in JSON format.'
)
elif self.format_type == data.FormatType.YAML:
system_message = (
'You are a helpful assistant that responds in YAML format.'
)
# Create the chat completion using the v1.x client API
response = self._client.chat.completions.create(
model=self.model_id,
messages=[
{'role': 'system', 'content': system_message},
{'role': 'user', 'content': prompt},
],
temperature=config.get('temperature', self.temperature),
max_tokens=config.get('max_output_tokens'),
top_p=config.get('top_p'),
n=1,
)
```
`config` is not passed by extract, so max_tokens is always `None`. If you step back a couple calls you find an arbitrary kwargs being processed that may contain `max_output_tokens`, but you can't pass into `.extract` without raising an error. I'll try to find a workaround but it'd be nice if max_tokens and other generation parameters could be passed into `.extract()`
Contributor guide
Assessment
This issue has not been assessed yet.