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.' ))
0 commit comments