mirror of
https://github.com/tcsenpai/suggestoor.git
synced 2025-06-04 18:30:19 +00:00
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
import os
|
|
import dotenv
|
|
from ollama import Client
|
|
from langchain_community.vectorstores import FAISS
|
|
from langchain_ollama import OllamaEmbeddings
|
|
|
|
dotenv.load_dotenv()
|
|
|
|
|
|
def load_vector_store():
|
|
embeddings = OllamaEmbeddings(
|
|
model=os.getenv("EMBEDDING_MODEL"),
|
|
base_url=os.getenv("OLLAMA_HOST"),
|
|
)
|
|
return FAISS.load_local(
|
|
"vector_store", embeddings, allow_dangerous_deserialization=True
|
|
)
|
|
|
|
|
|
def ollama_client():
|
|
return Client(
|
|
host=os.getenv("OLLAMA_HOST"),
|
|
)
|
|
|
|
|
|
def ollama_chat(client: Client, prompt: str):
|
|
db = load_vector_store()
|
|
|
|
# Get configurable parameters from environment
|
|
k_similar = int(
|
|
os.getenv("SIMILARITY_K", "5")
|
|
) # Number of similar chunks to retrieve
|
|
|
|
# Search for relevant context
|
|
relevant_chunks = db.similarity_search(prompt, k=k_similar)
|
|
|
|
context = "\n\n".join([doc.page_content for doc in relevant_chunks])
|
|
|
|
return client.chat(
|
|
model=os.getenv("OLLAMA_MODEL"),
|
|
messages=[
|
|
{
|
|
"role": "system",
|
|
"content": """You are a command-line tool expert. Your role is to suggest the most appropriate command-line tools based on the user's needs.
|
|
|
|
Rules:
|
|
1. ONLY suggest tools that are explicitly mentioned in the provided context
|
|
2. Start with the most fundamental/common tool for the task
|
|
3. Format your response as:
|
|
Primary tool: [tool name] - [brief description]
|
|
Usage: [basic usage example]
|
|
|
|
Related tools:
|
|
- [alternative tool] - [why/when to use it]
|
|
|
|
4. If no suitable tool is found in the context, say "I don't see a suitable tool in my current knowledge base for that specific task"
|
|
5. Be concise and direct
|
|
6. Include specific, relevant details from the man pages
|
|
7. Don't make assumptions about tools that aren't in the context
|
|
|
|
Here is the relevant documentation context:
|
|
{}
|
|
""".format(
|
|
context
|
|
),
|
|
},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
)
|