Skip to content

Commit

Permalink
WIP: Mirror Ollama chat API
Browse files Browse the repository at this point in the history
  • Loading branch information
TilmanGriesel committed Jan 29, 2025
1 parent 9b56b1c commit 648a38d
Showing 1 changed file with 10 additions and 22 deletions.
32 changes: 10 additions & 22 deletions services/api/src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@
from pathlib import Path

import elasticsearch
from core.ollama_proxy import OllamaProxy
from core.pipeline_config import ModelProvider, QueryPipelineConfig
from core.rag_pipeline import RAGQueryPipeline
from core.ollama_proxy import OllamaProxy
from dotenv import load_dotenv
from flask import Flask, Response, abort, jsonify, request, stream_with_context
from flask_limiter import Limiter
Expand Down Expand Up @@ -110,6 +110,7 @@ def load_systemprompt(base_path: str) -> str:

system_prompt_value = load_systemprompt(os.getenv("SYSTEM_PROMPT_PATH", os.getcwd()))


def get_env_param(param_name, converter=None, default=None):
value = os.getenv(param_name)
if value is None:
Expand Down Expand Up @@ -262,9 +263,11 @@ def tags():
def pull():
return proxy.pull()


ollama_proxy = OllamaProxy(os.getenv("OLLAMA_URL", "http://localhost:11434"))
register_ollama_routes(app, ollama_proxy)


@app.before_request
def before_request():
logger.info(f"Request {request.method} {request.path} from {request.remote_addr}")
Expand Down Expand Up @@ -441,19 +444,15 @@ def generate():


def handle_standard_response(
config: QueryPipelineConfig,
query: str,
conversation: list
config: QueryPipelineConfig, query: str, conversation: list
) -> Response:
rag = RAGQueryPipeline(config=config)

success = True
result = None
try:
result = rag.run_query(
query=query,
conversation=conversation,
print_response=False
query=query, conversation=conversation, print_response=False
)

if result:
Expand All @@ -474,30 +473,18 @@ def handle_standard_response(
"prompt_eval_count": 26,
"prompt_eval_duration": 342546000,
"eval_count": 282,
"eval_duration": 4535599000
"eval_duration": 4535599000,
}

except Exception as e:
success = False
logger.error(f"Error in RAG pipeline: {e}", exc_info=True)

if success and result:
latest_message = {
"role": "assistant",
"content": result["llm"]["replies"][0],
"timestamp": datetime.now().isoformat(),
}
conversation.append(latest_message)

return jsonify(
{
"success": success,
"timestamp": datetime.now().isoformat(),
"result": result,
"messages": conversation,
}
{"success": success, "timestamp": datetime.now().isoformat(), "result": result}
)


@app.route("/", methods=["GET"])
@app.route("/health", methods=["GET"])
def health_check():
Expand All @@ -511,6 +498,7 @@ def health_check():
}
)


@app.errorhandler(404)
def not_found_error(error):
return "", 404
Expand Down

0 comments on commit 648a38d

Please sign in to comment.