Custom Local Model Client (LM): dspy.Predict Outputs only the first character from the API response

Hi,

I have created a custom local model client (LM) as described in the documentation to connect with AI/ML APIs. And while the output from the API and from the class is as expected, upon processing through dspy.Predict or dspy.ChainOfThought it only returns the first character from the "answer" object.

Here is the script with the custom class:

import requests
import os
import dspy
from dotenv import load_dotenv
from dsp.modules.lm import LM

class AIMLAPI(LM):
    def __init__(self, model, api_key, **kwargs):
        super().__init__(model)
        self.api_key = api_key
        self.base_url = "https://api.aimlapi.com/chat/completions"
        self.kwargs.update(kwargs)
        self.provider = "AIMLAPI"

    def basic_request(self, prompt: str, **kwargs):
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}',
        }

        data = {
            "model": self.kwargs["model"],
            "messages": [{"role": "user", "content": prompt}],
            **{k: v for k, v in kwargs.items() if k in self.kwargs}
        }

        response = requests.post(self.base_url, headers=headers, json=data)
        response.raise_for_status()
        response_json = response.json()
        # Debugging 1: Print Json response from API: Good Response
        print("API Response:", response_json)

        self.history.append({
            "prompt": prompt,
            "response": response_json,
            "kwargs": kwargs,
        })

        return response_json

    def __call__(self, prompt, only_completed=True, return_sorted=False, **kwargs):
        response = self.basic_request(prompt, **kwargs)
        if 'choices' in response and len(response['choices']) > 0:
            return response['choices'][0]['message']['content']
        return "No valid response found"

# Load environment variables
load_dotenv()

# Set your API key and endpoint
AIML_API_KEY = os.getenv("AIML_API_KEY")
model = "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"

llama_lm = AIMLAPI(
    model=model,
    api_key=AIML_API_KEY,
    temperature=0.7,
    max_tokens=100,
    n=1
)

dspy.settings.configure(lm=llama_lm, temperature=0.7, max_tokens=1000, n=1)

prompt = "I want to ask about the universe"

# Debugging 2: calling the language model for the response: Good response
response = llama_lm(prompt)
print(f"Full Model Response: {response}")

# Process with Predict
processed_response = dspy.Predict("question -> answer", n=1)(question=prompt, lm=llama_lm)

# Debugging 3: Print the full processed response: Bad response [returning the first character]
print("Full Response from dspy.Predict:", processed_response)

# Attempt to extract answer from the processed response
try:
    # Inspect attributes and data
    attributes = dir(processed_response)
    print("Processed Response Attributes:", attributes)

    # Check if the answer is in the response
    if hasattr(processed_response, 'answer'):
        #Debugging 4: Print extracted answer from dspy.Predict: Bad response  [returning the first character]
        print(f"Extracted Answer: {processed_response.answer}")
    else:
        print("Answer attribute not found. Full response content:", processed_response)

except AttributeError as e:
    print(f"Error accessing answer: {e}. Full response content:", processed_response)

And here is the output of 4 print statements:

(venv) PS C:\Users\amagd\projects\falcon> python aimlapi_test.py

API Response: {'id': '8b54dfae3b597616-SEA', 'object': 'chat.completion', 'created': 1724015251, 'model': 'meta-llama/meta-llama-3.1-70b-instruct-turbo', 'prompt': [], 'choices': [{'finish_reason': 'eos', 'seed': 12916370056455328000, 'logprobs': None, 'index': 0, 'message': {'role': 'assistant', 'content': "What a vast and fascinating topic! The universe is full of mysteries and wonders. What would you like to know about the universe? Are you interested in:\n\n* The Big Bang and the origins of the universe?\n* Black holes and dark matter?\n* The search for extraterrestrial life?\n* The life cycle of stars and galaxies?\n* The mysteries of dark energy and the accelerating expansion of the universe?\n* Something else entirely?\n\nFeel free to ask me any question, and I'll do my best to provide an answer!"}}], 'usage': {'prompt_tokens': 40, 'completion_tokens': 243, 'total_tokens': 283}}

Full Model Response: What a vast and fascinating topic! The universe is full of mysteries and wonders. What would you like to know about the universe? Are you interested in:

* The Big Bang and the origins of the universe?
* Black holes and dark matter?
* The search for extraterrestrial life?
* The life cycle of stars and galaxies?
* The mysteries of dark energy and the accelerating expansion of the universe?
* Something else entirely?

Feel free to ask me any question, and I'll do my best to provide an answer!

API Response: {'id': '8b54dfbb1db27616-SEA', 'object': 'chat.completion', 'created': 1724015252, 'model': 'meta-llama/meta-llama-3.1-70b-instruct-turbo', 'prompt': [], 'choices': [{'finish_reason': 'eos', 'seed': 5768321966100472000, 'logprobs': None, 'index': 0, 'message': {'role': 'assistant', 'content': 'What would you like to know about the universe?'}}], 'usage': {'prompt_tokens': 120, 'completion_tokens': 25, 'total_tokens': 145}}

Full Response from dspy.Predict: Prediction(
    answer='W',
    completions=Completions(...)
) (38 completions omitted)

Processed Response Attributes: ['__class__', '__contains__', '__delattr__', '__delitem__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattr__', '__getattribute__', '__getitem__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__iter__', '__le__', '__len__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setitem__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_completions', '_store', 'completions', 'copy', 'from_completions', 'get', 'inputs', 'items', 'keys', 'labels', 'toDict', 'values', 'with_inputs', 'without']
Extracted Answer: W

Any idea how to get a proper response from dspy.Predict with this model?

stanfordnlp / dspy

Custom Local Model Client (LM): dspy.Predict Outputs only the first character from the API response #1407