Skip to content

Commit

Permalink
Remove unnecessary _call_api() abstraction
Browse files Browse the repository at this point in the history
  • Loading branch information
smokestacklightnin committed Apr 10, 2024
1 parent 1ce1982 commit fd5c34b
Showing 1 changed file with 2 additions and 14 deletions.
16 changes: 2 additions & 14 deletions ragna/assistants/_ollama.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,12 +62,8 @@ async def _assert_api_call_is_success(self, response: Response) -> None:
response_content=content,
)

async def _call_api(
self,
prompt: str,
sources: list[Source],
*,
max_new_tokens: int,
async def answer(
self, prompt: str, sources: list[Source], *, max_new_tokens: int = 256
) -> AsyncIterator[str]:
async with self._client.stream(
"POST",
Expand Down Expand Up @@ -103,14 +99,6 @@ async def _call_api(
if not json_data["done"]:
yield cast(str, json_data["message"]["content"])

async def answer(
self, prompt: str, sources: list[Source], *, max_new_tokens: int = 256
) -> AsyncIterator[str]:
async for chunk in self._call_api( # type: ignore[attr-defined, misc]
prompt, sources, max_new_tokens=max_new_tokens
):
yield chunk


class OllamaGemma2B(OllamaApiAssistant):
_MODEL = "gemma:2b"

0 comments on commit fd5c34b

Please sign in to comment.