aboutsummaryrefslogtreecommitdiff
path: root/autogpts/autogpt/tests/integration/memory/utils.py
blob: aea12832fc2b70497ac42ef33a18ce3fed3000d7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import numpy
import pytest
from pytest_mock import MockerFixture

import autogpt.memory.vector.memory_item as vector_memory_item
import autogpt.memory.vector.providers.base as memory_provider_base
from autogpt.config.config import Config
from autogpt.core.resource.model_providers import OPEN_AI_EMBEDDING_MODELS
from autogpt.memory.vector import get_memory
from autogpt.memory.vector.utils import Embedding


@pytest.fixture
def embedding_dimension(config: Config):
    return OPEN_AI_EMBEDDING_MODELS[config.embedding_model].embedding_dimensions


@pytest.fixture
def mock_embedding(embedding_dimension: int) -> Embedding:
    return numpy.full((1, embedding_dimension), 0.0255, numpy.float32)[0]


@pytest.fixture
def mock_get_embedding(mocker: MockerFixture, mock_embedding: Embedding):
    mocker.patch.object(
        vector_memory_item,
        "get_embedding",
        return_value=mock_embedding,
    )
    mocker.patch.object(
        memory_provider_base,
        "get_embedding",
        return_value=mock_embedding,
    )


@pytest.fixture
def memory_none(agent_test_config: Config, mock_get_embedding):
    was_memory_backend = agent_test_config.memory_backend

    agent_test_config.memory_backend = "no_memory"
    yield get_memory(agent_test_config)

    agent_test_config.memory_backend = was_memory_backend