-
-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
3 changed files
with
53 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,3 @@ | ||
from .ollama import Ollama | ||
from .openai import OpenAI | ||
from .localai import LocalAI |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,49 @@ | ||
from collections.abc import AsyncIterator | ||
|
||
from openai import APIError | ||
from openai import AsyncOpenAI as OpenAIClient | ||
|
||
from ..config import Setting | ||
from ..provider import Provider, ProviderError | ||
|
||
|
||
class LocalAI(Provider): | ||
name = "LocalAI" | ||
|
||
host = Setting(default="localhost") | ||
port = Setting(default=8080) | ||
model = Setting(default="lunademo") | ||
system_prompt = Setting( | ||
default=( | ||
"Based on the following user description, generate a corresponding Bash command. Focus solely " | ||
"on interpreting the requirements and translating them into a single, executable Bash command. " | ||
"Ensure accuracy and relevance to the user's description. The output should be a valid Bash " | ||
"command that directly aligns with the user's intent, ready for execution in a command-line " | ||
"environment. Output nothing except for the command. No code block, no English explanation, " | ||
"no start/end tags." | ||
) | ||
) | ||
|
||
@property | ||
def endpoint(self) -> str: | ||
# computed property because python descriptors need to be bound to an instance before access | ||
return f"http://{self.host}:{self.port}/api/generate" | ||
|
||
def __init__(self): | ||
self.client = OpenAIClient(base_url=self.endpoint) | ||
|
||
async def generate(self, prompt: str) -> AsyncIterator[str]: | ||
try: | ||
stream = await self.client.chat.completions.create( | ||
model=self.model, | ||
messages=[ | ||
{"role": "system", "content": self.system_prompt}, | ||
{"role": "user", "content": prompt} | ||
], | ||
stream=True, | ||
) | ||
async for chunk in stream: | ||
if chunk.choices[0].delta.content is not None: | ||
yield chunk.choices[0].delta.content | ||
except APIError as e: | ||
raise ProviderError(f"Something went wrong while querying OpenAI: {e}") from e |