Make API URL customizable

Quansight · Mar 31, 2024 · 39d2e44 · 39d2e44
1 parent 5f0ce9c
commit 39d2e44
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/ragna/assistants/_ollama.py b/ragna/assistants/_ollama.py
@@ -30,11 +30,16 @@ def _make_system_content(self, sources: list[Source]) -> str:
         return instruction + "\n\n".join(source.content for source in sources)
 
     async def _call_api(
-        self, prompt: str, sources: list[Source], *, max_new_tokens: int
+        self,
+        prompt: str,
+        sources: list[Source],
+        *,
+        max_new_tokens: int,
+        api_url: str = "http://localhost:11434/api/chat",
     ) -> AsyncIterator[str]:
         async with self._client.stream(
             "POST",
-            "http://localhost:11434/api/chat",  # TODO: Make this url customizable
+            api_url,
             headers={
                 "Content-Type": "application/json",
             },