From c496bb485c26106eecefa48b2c1be8c0bae2ecd1 Mon Sep 17 00:00:00 2001 From: Vasilije <8619304+Vasilije1990@users.noreply.github.com> Date: Fri, 28 Feb 2025 11:15:12 -0800 Subject: [PATCH] feat: Draft ollama test (#566) ## Description ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin ## Summary by CodeRabbit - **Tests** - Introduced new automated testing workflows for Ollama and Gemini, triggered by pull requests and manual dispatch. - The Ollama workflow sets up the service and executes a simple example test to enhance continuous integration. - Enhanced dependency update workflow with new triggers for push and pull request events, and added an optional debug logging parameter. - Added new capabilities for audio and image transcription within the Ollama API adapter. --------- Co-authored-by: Daniel Molnar --- .github/workflows/test_gemini.yml | 33 ++++++ .github/workflows/test_ollama.yml | 116 ++++++++++++++++++++ .github/workflows/upgrade_deps.yml | 23 +++- cognee/infrastructure/llm/ollama/adapter.py | 53 +++++++++ 4 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/test_gemini.yml create mode 100644 .github/workflows/test_ollama.yml diff --git a/.github/workflows/test_gemini.yml b/.github/workflows/test_gemini.yml new file mode 100644 index 000000000..e5bf0520f --- /dev/null +++ b/.github/workflows/test_gemini.yml @@ -0,0 +1,33 @@ +name: test | gemini + +on: + workflow_dispatch: + pull_request: + types: [labeled, synchronize] + + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + run_simple_example_test: + uses: ./.github/workflows/reusable_python_example.yml + with: + example-location: ./examples/python/simple_example.py + secrets: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }} + GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }} + EMBEDDING_PROVIDER: "gemini" + EMBEDDING_API_KEY: ${{ secrets.GEMINI_API_KEY }} + EMBEDDING_MODEL: "gemini/text-embedding-004" + EMBEDDING_ENDPOINT: "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004" + EMBEDDING_API_VERSION: "v1beta" + EMBEDDING_DIMENSIONS: 768 + EMBEDDING_MAX_TOKENS: 8076 + LLM_PROVIDER: "gemini" + LLM_API_KEY: ${{ secrets.GEMINI_API_KEY }} + LLM_MODEL: "gemini/gemini-1.5-flash" + LLM_ENDPOINT: "https://generativelanguage.googleapis.com/" + LLM_API_VERSION: "v1beta" diff --git a/.github/workflows/test_ollama.yml b/.github/workflows/test_ollama.yml new file mode 100644 index 000000000..030d6e3e1 --- /dev/null +++ b/.github/workflows/test_ollama.yml @@ -0,0 +1,116 @@ +name: test | ollama + +on: + workflow_dispatch: + pull_request: + types: [ labeled, synchronize ] + +jobs: + + run_simple_example_test: + + # needs 16 Gb RAM for phi4 + runs-on: buildjet-4vcpu-ubuntu-2204 +# services: +# ollama: +# image: ollama/ollama +# ports: +# - 11434:11434 + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: '3.12.x' + + - name: Install Poetry + uses: snok/install-poetry@v1.4.1 + with: + virtualenvs-create: true + virtualenvs-in-project: true + installer-parallel: true + + - name: Install dependencies + run: | + poetry install --no-interaction --all-extras + poetry add torch + +# - name: Install ollama +# run: curl -fsSL https://ollama.com/install.sh | sh +# - name: Run ollama +# run: | +# ollama serve --openai & +# ollama pull llama3.2 & +# ollama pull avr/sfr-embedding-mistral:latest + + - name: Start Ollama container + run: | + docker run -d --name ollama -p 11434:11434 ollama/ollama + sleep 5 + docker exec -d ollama bash -c "ollama serve --openai" + + - name: Check Ollama logs + run: docker logs ollama + + - name: Wait for Ollama to be ready + run: | + for i in {1..30}; do + if curl -s http://localhost:11434/v1/models > /dev/null; then + echo "Ollama is ready" + exit 0 + fi + echo "Waiting for Ollama... attempt $i" + sleep 2 + done + echo "Ollama failed to start" + exit 1 + + - name: Pull required Ollama models + run: | + curl -X POST http://localhost:11434/api/pull -d '{"name": "phi4"}' + curl -X POST http://localhost:11434/api/pull -d '{"name": "avr/sfr-embedding-mistral:latest"}' + + - name: Call ollama API + run: | + curl -X POST http://localhost:11434/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "phi4", + "stream": false, + "messages": [ + { "role": "system", "content": "You are a helpful assistant." }, + { "role": "user", "content": "Whatever I say, answer with Yes." } + ] + }' + curl -X POST http://127.0.0.1:11434/v1/embeddings \ + -H "Content-Type: application/json" \ + -d '{ + "model": "avr/sfr-embedding-mistral:latest", + "input": "This is a test sentence to generate an embedding." + }' + + - name: Dump Docker logs + run: | + docker ps + docker logs $(docker ps --filter "ancestor=ollama/ollama" --format "{{.ID}}") + + + - name: Run example test + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }} + GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }} + PYTHONFAULTHANDLER: 1 + LLM_PROVIDER: "ollama" + LLM_API_KEY: "ollama" + LLM_ENDPOINT: "http://localhost:11434/v1/" + LLM_MODEL: "phi4" + EMBEDDING_PROVIDER: "ollama" + EMBEDDING_MODEL: "avr/sfr-embedding-mistral:latest" + EMBEDDING_ENDPOINT: "http://localhost:11434/v1/" + EMBEDDING_DIMENSIONS: "4096" + HUGGINGFACE_TOKENIZER: "Salesforce/SFR-Embedding-Mistral" + run: poetry run python ./examples/python/simple_example.py diff --git a/.github/workflows/upgrade_deps.yml b/.github/workflows/upgrade_deps.yml index 934a3ed6a..327ee08cc 100644 --- a/.github/workflows/upgrade_deps.yml +++ b/.github/workflows/upgrade_deps.yml @@ -2,8 +2,29 @@ name: Update Poetry Dependencies on: schedule: - - cron: '0 3 * * 0' + - cron: '0 3 * * 0' # Runs at 3 AM every Sunday + push: + paths: + - 'poetry.lock' + - 'pyproject.toml' + branches: + - main + - dev + pull_request: + paths: + - 'poetry.lock' + - 'pyproject.toml' + types: [opened, synchronize, reopened] + branches: + - main + - dev workflow_dispatch: + inputs: + debug_enabled: + type: boolean + description: 'Run the update with debug logging' + required: false + default: false jobs: update-dependencies: diff --git a/cognee/infrastructure/llm/ollama/adapter.py b/cognee/infrastructure/llm/ollama/adapter.py index 4eb392739..ac6b91af4 100644 --- a/cognee/infrastructure/llm/ollama/adapter.py +++ b/cognee/infrastructure/llm/ollama/adapter.py @@ -4,6 +4,8 @@ import instructor from cognee.infrastructure.llm.llm_interface import LLMInterface from cognee.infrastructure.llm.config import get_llm_config from openai import OpenAI +import base64 +import os class OllamaAPIAdapter(LLMInterface): @@ -42,3 +44,54 @@ class OllamaAPIAdapter(LLMInterface): ) return response + + def create_transcript(self, input_file: str) -> str: + """Generate an audio transcript from a user query.""" + + if not os.path.isfile(input_file): + raise FileNotFoundError(f"The file {input_file} does not exist.") + + with open(input_file, "rb") as audio_file: + transcription = self.aclient.audio.transcriptions.create( + model="whisper-1", # Ensure the correct model for transcription + file=audio_file, + language="en", + ) + + # Ensure the response contains a valid transcript + if not hasattr(transcription, "text"): + raise ValueError("Transcription failed. No text returned.") + + return transcription.text + + def transcribe_image(self, input_file: str) -> str: + """Transcribe content from an image using base64 encoding.""" + + if not os.path.isfile(input_file): + raise FileNotFoundError(f"The file {input_file} does not exist.") + + with open(input_file, "rb") as image_file: + encoded_image = base64.b64encode(image_file.read()).decode("utf-8") + + response = self.aclient.chat.completions.create( + model=self.model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What’s in this image?"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}, + }, + ], + } + ], + max_tokens=300, + ) + + # Ensure response is valid before accessing .choices[0].message.content + if not hasattr(response, "choices") or not response.choices: + raise ValueError("Image transcription failed. No response received.") + + return response.choices[0].message.content