Skip to content

Commit 1d10a22

Browse files
committed
Starting migration of knowledge items
1 parent defa36b commit 1d10a22

4 files changed

Lines changed: 120 additions & 21 deletions

File tree

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
from django.core.management.base import BaseCommand
2+
from django.core.files.base import ContentFile
3+
from back.apps.language_model.models import KnowledgeBase, KnowledgeItem, KnowledgeItemImage
4+
import requests
5+
6+
7+
class Command(BaseCommand):
8+
help = 'Migrate knowledge items and images from a source deployment to a destination deployment'
9+
# command: python manage.py migrate_knowledge_items http://source-deployment.com source_kb_name destination_kb_name --token=token
10+
11+
def add_arguments(self, parser):
12+
parser.add_argument('source_base_url', type=str, help='Base URL of the source deployment')
13+
parser.add_argument('source_base_name', type=str, help='Name of the source knowledge base')
14+
parser.add_argument('destination_base_name', type=str, help='Name of the destination knowledge base')
15+
parser.add_argument('--token', type=str, help='Token for header authentication')
16+
17+
def handle(self, *args, **options):
18+
source_back_url = options['source_base_url']
19+
source_kb_name = options['source_base_name']
20+
destination_kb_name = options['destination_base_name']
21+
token = options['token']
22+
header = {'Authorization': f'Token {token}'} if token else {}
23+
24+
# Retrieve the destination knowledge base by name
25+
try:
26+
destination_base = KnowledgeBase.objects.get(name=destination_kb_name)
27+
except KnowledgeBase.DoesNotExist:
28+
self.stdout.write(self.style.ERROR(f'No knowledge base found with name "{destination_kb_name}" in the destination deployment.'))
29+
return
30+
31+
# Retrieve the source knowledge base by name to check if it exists
32+
response = requests.get(f'{source_back_url}/api/language-model/knowledge-bases/?name={source_kb_name}', headers=header)
33+
source_base_data = response.json()
34+
35+
if not source_base_data:
36+
self.stdout.write(self.style.ERROR(f'No knowledge base found with name "{source_kb_name}" in the source deployment.'))
37+
return
38+
39+
# Retrieve knowledge items from the source deployment for the specified knowledge base
40+
response = requests.get(f'{source_back_url}/api/language-model/knowledge-items/?knowledge_base__name={source_kb_name}', headers=header)
41+
knowledge_items_data = response.json()
42+
43+
if knowledge_items_data['count'] == 0:
44+
self.stdout.write(self.style.ERROR(f'No knowledge items found for knowledge base "{source_kb_name}" in the source deployment.'))
45+
return
46+
47+
for item_data in knowledge_items_data:
48+
# Create a new knowledge item in the destination deployment
49+
knowledge_item = KnowledgeItem(
50+
knowledge_base=destination_base,
51+
title=item_data['title'],
52+
content=item_data['content'],
53+
url=item_data['url'],
54+
section=item_data['section'],
55+
role=item_data['role'],
56+
page_number=item_data['page_number'],
57+
metadata=item_data['metadata']
58+
)
59+
knowledge_item.save()
60+
61+
# # Retrieve and create associated images
62+
# images_data = item_data['images']
63+
# for image_data in images_data:
64+
# image_response = requests.get(image_data['image_url'])
65+
# image_content = ContentFile(image_response.content)
66+
# knowledge_item_image = KnowledgeItemImage(
67+
# knowledge_item=knowledge_item,
68+
# image_caption=image_data['image_caption']
69+
# )
70+
# knowledge_item_image.image_file.save(image_data['image_name'], image_content)
71+
# knowledge_item_image.save()
72+
73+
self.stdout.write(self.style.SUCCESS('Knowledge items and images migrated successfully.'))

back/back/apps/language_model/views/data.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,20 @@
2525
)
2626

2727

28+
class KnowledgeBaseFilter(django_filters.FilterSet):
29+
class Meta:
30+
model = KnowledgeBase
31+
fields = {
32+
'name': ['exact'],
33+
}
34+
35+
2836
class KnowledgeBaseAPIViewSet(viewsets.ModelViewSet):
2937
queryset = KnowledgeBase.objects.all()
3038
serializer_class = KnowledgeBaseSerializer
3139
filter_backends = [DjangoFilterBackend, SearchFilter, OrderingFilter]
3240
search_fields = ['name']
41+
filterset_class = KnowledgeBaseFilter
3342

3443
def get_queryset(self):
3544
if self.kwargs.get("pk"):

0 commit comments

Comments
 (0)