Skip to content

Commit

Permalink
update interface.py
Browse files Browse the repository at this point in the history
  • Loading branch information
itsmvd committed Dec 23, 2024
1 parent 8c98962 commit 8eaad30
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions timesketch/lib/llms/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
"""Interface for LLM providers."""

import string
from typing import Optional

from flask import current_app

Expand Down Expand Up @@ -82,13 +83,13 @@ def prompt_from_template(self, template: str, kwargs: dict) -> str:
formatter = string.Formatter()
return formatter.format(template, **kwargs)

def generate(self, prompt: str) -> str:
def generate(self, prompt: str, response_schema: Optional[dict] = None) -> str:
"""Generate a response from the LLM provider.
Args:
prompt: The prompt to generate a response for.
temperature: The temperature to use for the response.
stream: Whether to stream the response.
response_schema: An optional JSON schema to define the expected
response format.
Returns:
The generated response.
Expand Down

0 comments on commit 8eaad30

Please sign in to comment.