add natural language queries

agent.py

@@ -1,13 +1,26 @@
-import transformers
+
+import re
+import json
+
+import functools
+import functions
 import torch
 from tools import tools
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
-    BitsAndBytesConfig,
-    TextIteratorStreamer,
+    BitsAndBytesConfig
 )
 
+# Get all functions from functions.py
+all_functions = [func for func in dir(functions) if callable(
+    getattr(functions, func)) and not func.startswith("__")]
+
+# Create names_to_function dict containing partials for all functions in functions.py
+names_to_functions = {func: functools.partial(
+    getattr(functions, func)) for func in all_functions}
+
+
 model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 
 # specify how to quantize the model
@@ -23,11 +36,35 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 
-messages = [
-    {"role": "system", "content": "You are a movie search assistant bot who uses TMDB to help users find movies. You should respond with movie IDs and natural language text summaries when asked for movie recommendations. You should only provide the movie ID and the summary, nothing else."},
-    {"role": "user", "content": "Can you recommend a good action movie?"},
-]
+def extract_function_call(output):
+    match = re.search(r'<\|python_tag\|>(.*)<\|eom_id\|>', output)
+    if match:
+        function_call = match.group(1)
+        return json.loads(function_call)
+    else:
+        return None
+
+
+def chatbot(query):
+    messages = [
+        {"role": "system", "content": "You are a movie search assistant bot who uses TMDB to help users find movies. Think step by step and identify the sequence of function calls that will help to answer."},
+        {"role": "user", "content": query},
+    ]
+
+    tokenized_chat = tokenizer.apply_chat_template(
+        messages, tools=tools, add_generation_prompt=True, tokenize=True, return_tensors="pt")
 
+    outputs = model.generate(tokenized_chat, max_new_tokens=128)
+    answer = tokenizer.batch_decode(outputs[:, tokenized_chat.shape[1]:])[0]
+    tool_call = extract_function_call(answer)
 
-inputs = tokenizer.apply_chat_template(
-    messages, tools=tools, add_generation_prompt=True)
+    if tool_call:
+        function_name = tool_call['name']
+        function_params = tool_call['parameters']
+        print("\nfunction_name: ", function_name,
+              "\nfunction_params: ", function_params)
+        function_result = names_to_functions[function_name](**function_params)
+        print(function_result['results'])
+        return function_result['results']
+    else:
+        print("No tool calls found in the answer.")

app.py

@@ -1,4 +1,5 @@
 from flask import Flask, render_template, request, jsonify
+from agent import chatbot
 from functions import query_tmdb
 
 app = Flask(__name__)
@@ -83,5 +84,16 @@ def cast_filter(cast_id):
                            cast_info=cast_info)
 
 
+@app.route('/query', methods=['GET', 'POST'])
+def query():
+    if request.method == 'POST':
+        query = request.form.get('query')
+        result = chatbot(query)  # Call the chatbot function from agent.py
+        print("Got result: ", result)
+        print("Movie ids are: ", [r['id'] for r in result])
+        return render_template('query.html', result=result)
+    return render_template('query.html')
+
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)

functions.py

@@ -14,47 +14,13 @@ def query_tmdb(endpoint, params={}):
     return response.json()
 
 
-def discover_movie(certification=None, certification_gte=None, certification_lte=None,
-                   certification_country=None, include_adult=False, include_video=False,
-                   language="en-US", page=1, primary_release_year=None,
-                   primary_release_date_gte=None, primary_release_date_lte=None,
-                   region=None, release_date_gte=None, release_date_lte=None,
-                   sort_by="popularity.desc", vote_average_gte=None, vote_average_lte=None):
+def discover_movie(include_adult=False, include_video=False, language="en-US", page=1, sort_by="popularity.desc", **kwargs):
     endpoint = f"{BASE_URL}/discover/movie"
-    params = {
-        'include_adult': include_adult,
-        'include_video': include_video,
-        'language': language,
-        'page': page,
-        'sort_by': sort_by
-    }
-
-    # Add optional parameters if they're provided
-    if certification:
-        params['certification'] = certification
-    if certification_gte:
-        params['certification.gte'] = certification_gte
-    if certification_lte:
-        params['certification.lte'] = certification_lte
-    if certification_country:
-        params['certification_country'] = certification_country
-    if primary_release_year:
-        params['primary_release_year'] = primary_release_year
-    if primary_release_date_gte:
-        params['primary_release_date.gte'] = primary_release_date_gte
-    if primary_release_date_lte:
-        params['primary_release_date.lte'] = primary_release_date_lte
-    if region:
-        params['region'] = region
-    if release_date_gte:
-        params['release_date.gte'] = release_date_gte
-    if release_date_lte:
-        params['release_date.lte'] = release_date_lte
-    if vote_average_gte:
-        params['vote_average.gte'] = vote_average_gte
-    if vote_average_lte:
-        params['vote_average.lte'] = vote_average_lte
+    params = {}
 
+    for key, value in kwargs.items():
+        if value is not None:
+            params[key] = value
     response = query_tmdb(endpoint, params=params)
     return response
 

templates/base.html

@@ -16,6 +16,7 @@
             <li><a href="/">Popular Movies</a></li>
             <li><a href="/genres">By Genre</a></li>
             <li><a href="/cast">By Cast</a></li>
+            <li><a href="/query">Query</a></li>
         </ul>
     </div>
     <div id="main-content">

templates/query.html

@@ -0,0 +1,23 @@
+{% extends 'base.html' %}
+
+{% block content %}
+<h1>Query</h1>
+<form action="/query" method="post">
+    <input type="text" name="query" placeholder="Enter your query...">
+    <input type="submit" value="Submit">
+</form>
+
+{% if result %}
+<h2>Query Result</h2>
+<div>
+    {% for movie in result %}
+    <div class="tile-element" onclick="showDetails({{ movie.id }})">
+        <img src="https://image.tmdb.org/t/p/w500{{ movie.poster_path }}" alt="{{ movie.title }}">
+        <h4>{{ movie.title }}</h4>
+    </div>
+    {% endfor %}
+</div>
+{% elif result is defined %}
+<p>No results found.</p>
+{% endif %}
+{% endblock %}