google / google/langextract

feat: Add max_tokens support to .extract(...)

Open
#72 1 comment 1 reaction 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.