mirror of
https://github.com/khoj-ai/khoj.git
synced 2024-11-23 23:48:56 +01:00
7dfbcd2e5a
- Ensure langchain less than 0.2.0 is used, to prevent breaking ChatOpenAI, PyMuPDF usage due to their deprecation after 0.2.0 - Set subscription renewal date to a timezone aware datetime - Use logger.warning instead of logger.warn as latter is deprecated - Use `model_dump' not deprecated dict to get all configured content_types
95 lines
2.3 KiB
Python
95 lines
2.3 KiB
Python
import os
|
|
from datetime import datetime
|
|
|
|
import factory
|
|
from django.utils.timezone import make_aware
|
|
|
|
from khoj.database.models import (
|
|
ChatModelOptions,
|
|
Conversation,
|
|
KhojApiUser,
|
|
KhojUser,
|
|
OfflineChatProcessorConversationConfig,
|
|
OpenAIProcessorConversationConfig,
|
|
SearchModelConfig,
|
|
Subscription,
|
|
UserConversationConfig,
|
|
)
|
|
|
|
|
|
class UserFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = KhojUser
|
|
|
|
username = factory.Faker("name")
|
|
email = factory.Faker("email")
|
|
password = factory.Faker("password")
|
|
uuid = factory.Faker("uuid4")
|
|
|
|
|
|
class ApiUserFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = KhojApiUser
|
|
|
|
user = None
|
|
name = factory.Faker("name")
|
|
token = factory.Faker("password")
|
|
|
|
|
|
class ChatModelOptionsFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = ChatModelOptions
|
|
|
|
max_prompt_size = 2000
|
|
tokenizer = None
|
|
chat_model = "mistral-7b-instruct-v0.1.Q4_0.gguf"
|
|
model_type = "offline"
|
|
|
|
|
|
class UserConversationProcessorConfigFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = UserConversationConfig
|
|
|
|
user = factory.SubFactory(UserFactory)
|
|
setting = factory.SubFactory(ChatModelOptionsFactory)
|
|
|
|
|
|
class OfflineChatProcessorConversationConfigFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = OfflineChatProcessorConversationConfig
|
|
|
|
enabled = True
|
|
|
|
|
|
class OpenAIProcessorConversationConfigFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = OpenAIProcessorConversationConfig
|
|
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
|
|
|
|
class ConversationFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = Conversation
|
|
|
|
user = factory.SubFactory(UserFactory)
|
|
|
|
|
|
class SearchModelFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = SearchModelConfig
|
|
|
|
name = "default"
|
|
model_type = "text"
|
|
bi_encoder = "thenlper/gte-small"
|
|
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
|
|
|
|
|
class SubscriptionFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = Subscription
|
|
|
|
user = factory.SubFactory(UserFactory)
|
|
type = "standard"
|
|
is_recurring = False
|
|
renewal_date = make_aware(datetime.strptime("2100-04-01", "%Y-%m-%d"))
|